From 2b9fbc5c3dc59e86187350ba875eee819e8f2bf0 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 5 May 2026 18:50:17 -0700 Subject: [PATCH 01/63] refactor nvte_get_fused_attn_backend with FE calls Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 516 ++++++------------ .../fused_attn_f16_arbitrary_seqlen.cu | 136 +++++ .../fused_attn_f16_arbitrary_seqlen.h | 22 + .../common/fused_attn/fused_attn_fp8.cu | 101 ++++ .../common/fused_attn/fused_attn_fp8.h | 25 + .../include/transformer_engine/fused_attn.h | 50 +- .../jax/csrc/extensions/attention.cpp | 32 +- .../pytorch/csrc/extensions/attention.cpp | 17 +- 8 files changed, 539 insertions(+), 360 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 141767b803..615f7c2a03 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -226,357 +226,189 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { } } +namespace { + +// Per-thread storage for the message string handed back through +// NVTEFusedAttnBackendStatus::message. Re-used (cleared + re-populated) on every call to +// nvte_get_fused_attn_backend on this thread, which is exactly the lifetime documented in the +// public header. +thread_local std::string g_fused_attn_backend_status_buffer; + +// Apply (code, msg) to *out_status (if non-null), routing the message through the +// thread-local buffer so the returned `const char*` outlives this function call. +void set_status(NVTEFusedAttnBackendStatus *out_status, cudnn_frontend::error_code_t code, + const std::string &message) { + if (out_status == nullptr) return; + g_fused_attn_backend_status_buffer = message; + out_status->code = static_cast(code); + out_status->message = g_fused_attn_backend_status_buffer.c_str(); +} + +void set_status(NVTEFusedAttnBackendStatus *out_status, const cudnn_frontend::error_t &err) { + set_status(out_status, err.code, err.err_msg); +} + +void set_ok(NVTEFusedAttnBackendStatus *out_status) { + set_status(out_status, cudnn_frontend::error_code_t::OK, ""); +} + +} // namespace + // select a backend for fused attention +// +// Routing flow: +// 1. Apply TE post-filters that encode policies cuDNN-FE doesn't model directly: +// a. requires_64bit_ragged_offset -> cudnn >= 9.5 +// b. qkv_format == THD requires a padding-style mask +// c. cuDNN <= 9.15 + is_training + bshd/sbhd + max_seqlen_kv % 128 != 0 + +// cuda_graph + non-padding mask is rejected (known capture quirk) +// 2. Dispatch by dtype to the appropriate probe(s): +// - FP8 (E4M3/E5M2): is_supported_fp8_fwd (+ is_supported_fp8_bwd if training) +// - FP16/BF16: is_supported_f16_fwd (+ is_supported_f16_bwd if training) +// The probes call the same _impl that the executor uses, with workspace=nullptr. +// They run validate -> build_operation_graph -> create_execution_plans -> +// check_support -> build_plans, and populate a thread-local cache that the +// executor cache-hits on. +// 3. Return the selected backend, or NVTE_No_Backend if any probe rejects. +// +// When `out_status` is non-null, it is filled with a code + message describing the +// rejection (or {OK, ""} on success). TE post-filter rejections synthesize an +// INVALID_VALUE entry; probe rejections forward the cuDNN-FE / NVTE_CHECK error verbatim. 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) { + bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, + NVTEScalingMode scaling_mode, 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, cudnnHandle_t handle, + NVTEFusedAttnBackendStatus *out_status) { 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); + // Initialize to OK so callers get a clean status on the success path without us having to + // remember to set it at every return. + set_ok(out_status); 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(); - // For ragged offsets we only support 32-bit prior to cuDNN 9.5 - // Only used when THD format is requested. + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const auto cudnn_runtime_version = cudnnGetVersion(); + + // ---------- TE post-filters (apply before delegating to cuDNN-FE) ---------- + + // (1) Ragged-offset width: cuDNN < 9.5 only supports 32-bit offsets. 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 && - // 8.9: t3hd, max_s=512, d=64, padding - ((cudnn_runtime_version >= 8900 && sm_arch_ < 100 && - qkv_layout == NVTE_QKV_Layout::NVTE_T3HD && max_seqlen_q == max_seqlen_kv && - max_seqlen_q <= 512 && head_dim_qk == 64 && head_dim_v == 64 && - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - // 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) { - if (cudnn_runtime_version >= 8900) { - backend = NVTE_Fused_Attn_Backend::NVTE_FP8; - } else { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: FP8 fused attention is supported by cuDNN 8.9.0+." - " Please upgrade your cuDNN version if possible." - << std::endl; - } - } else if ((q_dtype == NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16)) { - bool flag_m512 = false; - bool flag_arb = false; - if ((sm_arch_ == 80 || sm_arch_ == 90) && (max_seqlen_q <= 512 && max_seqlen_q % 64 == 0) && - (max_seqlen_kv <= 512 && max_seqlen_kv % 64 == 0) && (head_dim_qk == 64) && - (head_dim_v == 64) && (num_attn_heads == num_gqa_groups) && - ((bias_type == NVTE_Bias_Type::NVTE_NO_BIAS) || - (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS)) && - ((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 && - max_seqlen_q == max_seqlen_kv) || - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) && - ((qkv_layout == NVTE_QKV_Layout::NVTE_SB3HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_SBHD_SB2HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_BS3HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_BSHD_BS2HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD)) && - ((window_size_left == -1) && (window_size_right == -1 || window_size_right == 0)) && - !requires_64bit_ragged_offset && - (softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) && !return_max_logit) { - flag_m512 = true; + if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { + set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, + "Configuration requires 64-bit ragged offsets, which require cuDNN >= 9.5."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + // (2) THD requires a padding-style mask. + if (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 && + attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { + set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, + "THD-format attention requires a padding-style mask " + "(PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT)."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + // (3) cuDNN-Graph capture quirk on cuDNN <= 9.15: training + bshd/sbhd with + // max_seqlen_kv % 128 != 0 + cuda_graph + non-padding mask hangs/miscompiles. + if (cudnn_runtime_version <= 91500 && 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) { + set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, + "Known cuDNN <= 9.15 capture quirk: training + bshd/sbhd + " + "max_seqlen_kv % 128 != 0 + cuda_graph + non-padding mask is unsupported."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + // ---------- Dispatch by dtype ---------- + + // Probes use a single-batch graph; capability checks in cuDNN-FE are batch-agnostic. + constexpr size_t probe_batch = 1; + // bottom_right_diagonal is a runtime API knob the router doesn't see; the BRCM-via-mask + // case is captured by attn_mask_type, so we probe with the default top-left alignment. + constexpr bool probe_bottom_right_diagonal = false; + + const bool is_fp8 = + (q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2); + const bool is_f16_or_bf16 = + (q_dtype == NVTEDType::kNVTEFloat16 || q_dtype == NVTEDType::kNVTEBFloat16); + + if (is_fp8) { + // TE-only FP8 post-filters: no 64-bit ragged offsets, no max-logit output. + if (requires_64bit_ragged_offset) { + set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, + "FP8 fused attention does not support 64-bit ragged offsets."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - 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_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)) && - // 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; + if (return_max_logit) { + set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, + "FP8 fused attention does not support return_max_logit."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (((max_seqlen_q > 512) || (max_seqlen_kv > 512)) && (flag_arb == true)) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; + const DType q_t = static_cast(q_dtype); + const DType o_t = static_cast(o_dtype); + auto fwd_status = is_supported_fp8_fwd( + probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, is_training, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, probe_bottom_right_diagonal, q_t, o_t, scaling_mode, + handle); + if (fwd_status.is_bad()) { + set_status(out_status, fwd_status); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if ((max_seqlen_q <= 512) && (max_seqlen_kv <= 512)) { - if (flag_arb == true) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; - } else if ((flag_arb == false) && (flag_m512 == true)) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen; - } - int env_backend = static_cast(backend); - env_backend = transformer_engine::getenv("NVTE_FUSED_ATTN_BACKEND", env_backend); - if (((env_backend == static_cast(NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen)) && - flag_m512) || - ((env_backend == static_cast(NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen)) && - flag_arb)) { - backend = static_cast(env_backend); + if (is_training) { + auto bwd_status = is_supported_fp8_bwd( + probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, + o_t, scaling_mode, handle); + if (bwd_status.is_bad()) { + set_status(out_status, bwd_status); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } - if (cudnn_runtime_version < 8901 && - backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: FP16/BF16 fused attention is supported by cuDNN 8.9.1+." - " Please upgrade your cuDNN version if possible." - << std::endl; - } - 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; - } - 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 ((cudnn_runtime_version <= 91500) && 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; + return NVTE_Fused_Attn_Backend::NVTE_FP8; + } + + if (is_f16_or_bf16) { + const DType q_t = static_cast(q_dtype); + auto fwd_status = is_supported_f16_fwd( + probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, + softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, q_t, + handle); + if (fwd_status.is_bad()) { + set_status(out_status, fwd_status); + 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; - } + if (is_training) { + auto bwd_status = is_supported_f16_bwd( + probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, + handle); + if (bwd_status.is_bad()) { + set_status(out_status, bwd_status); + 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_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, + "Unsupported Q dtype for fused attention " + "(only FP16/BF16/FP8_E4M3/FP8_E5M2 are routable)."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } // NVTE fused attention FWD with separate Q, K and V @@ -661,11 +493,14 @@ 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); + const NVTEDType O_type = static_cast(output_O->data.dtype); + const NVTEScalingMode scaling_mode = input_Q->scaling_mode; 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); + is_training, Q_type, KV_type, O_type, scaling_mode, 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, /*deterministic=*/false, handle, + /*out_status=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, is_training, attn_scale, @@ -747,11 +582,14 @@ 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); + const NVTEDType O_type = static_cast(input_O->data.dtype); + const NVTEScalingMode scaling_mode = input_Q->scaling_mode; 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); + /*is_training=*/true, Q_type, KV_type, O_type, scaling_mode, 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=*/false, cuda_graph, deterministic, + handle, /*out_status=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); 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..57ca14a3e1 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 @@ -1333,4 +1333,140 @@ void fused_attn_arbitrary_seqlen_bwd( NVTE_ERROR("Unexpected workspace_size."); } } + +namespace { +// Probe-time defaults for runtime-only quantities the router doesn't see (paged-KV dims, +// ragged max-tokens, bias dims). These produce a graph whose support surface matches the +// real executor's: for non-paged / non-ragged paths these are unused inside the impl; +// for ragged-THD we rebind to worst-case bounds; for paged we use 1 page of full s_kv per +// batch (= same dims as non-paged), so cuDNN-FE applies the paged-attention support rules. +struct ProbeDims { + 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; +}; + +ProbeDims compute_probe_dims(int64_t batch, int64_t num_attn_heads, int64_t max_seqlen_q, + int64_t max_seqlen_kv, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type) { + const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(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 NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + const bool has_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + + ProbeDims d{}; + d.max_b = (is_ragged_q || is_ragged_kv) ? batch : 0; + d.max_t_q = is_ragged_q ? batch * max_seqlen_q : 0; + d.max_t_kv = is_ragged_kv ? batch * max_seqlen_kv : 0; + d.num_pages_k = is_paged_kv ? batch : 0; + d.num_pages_v = is_paged_kv ? batch : 0; + d.page_size_k = is_paged_kv ? max_seqlen_kv : 0; + d.page_size_v = is_paged_kv ? max_seqlen_kv : 0; + d.max_pages_per_seq_k = is_paged_kv ? 1 : 0; + d.max_pages_per_seq_v = is_paged_kv ? 1 : 0; + d.bias_b = has_bias ? batch : 0; + d.bias_h = has_bias ? num_attn_heads : 0; + d.bias_sq = has_bias ? max_seqlen_q : 0; + d.bias_skv = has_bias ? max_seqlen_kv : 0; + return d; +} +} // namespace + +cudnn_frontend::error_t is_supported_f16_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, + bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, DType q_dtype, cudnnHandle_t handle) { + const ProbeDims d = compute_probe_dims(static_cast(batch), + static_cast(num_attn_heads), + static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), qkv_layout, + bias_type); + const NVTE_QKV_Format o_format = nvte_get_q_format(qkv_layout); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( + static_cast(batch), static_cast(num_attn_heads), + static_cast(num_gqa_groups), static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), static_cast(head_dim_qk), + static_cast(head_dim_v), d.max_b, d.max_t_q, d.max_t_kv, d.num_pages_k, + d.num_pages_v, d.page_size_k, d.page_size_v, d.max_pages_per_seq_k, + d.max_pages_per_seq_v, d.bias_b, d.bias_h, d.bias_sq, d.bias_skv, is_training, + return_max_logit, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, bias_type, + mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + /*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, + get_cudnn_fe_dtype(q_dtype), /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return {cudnn_frontend::error_code_t::OK, ""}; + } catch (const std::exception &e) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + } catch (...) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, + "is_supported_f16_fwd: unknown failure"}; + } +} + +cudnn_frontend::error_t is_supported_f16_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType q_dtype, cudnnHandle_t handle) { + const ProbeDims d = compute_probe_dims(static_cast(batch), + static_cast(num_attn_heads), + static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), qkv_layout, + bias_type); + const NVTE_QKV_Format o_format = nvte_get_q_format(qkv_layout); + const NVTE_QKV_Format do_format = o_format; + const NVTE_QKV_Layout dqkv_layout = qkv_layout; + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( + static_cast(batch), static_cast(num_attn_heads), + static_cast(num_gqa_groups), static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), static_cast(head_dim_qk), + static_cast(head_dim_v), d.max_b, d.max_t_q, d.max_t_kv, d.bias_b, d.bias_h, + d.bias_sq, d.bias_skv, /*scaling_factor=*/1.0f, 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=*/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, + get_cudnn_fe_dtype(q_dtype), /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return {cudnn_frontend::error_code_t::OK, ""}; + } catch (const std::exception &e) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + } catch (...) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, + "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..38cf48c1f0 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 @@ -12,6 +12,7 @@ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ #include +#include #include "common/common.h" #include "transformer_engine/fused_attn.h" @@ -47,6 +48,27 @@ void fused_attn_arbitrary_seqlen_bwd( const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); +// Probe: drives cuDNN-FE (validate -> build_operation_graph -> create_execution_plans -> +// check_support -> build_plans) for an F16/BF16 forward graph with the given configuration. +// Returns the cuDNN-FE status: error_code_t::OK iff the graph compiles end-to-end. On OK, +// the built graph is inserted into the same thread-local cache used by +// fused_attn_arbitrary_seqlen_fwd_impl, so the executor cache-hits on matching descriptors. +// On rejection, err_msg contains the underlying cuDNN-FE / NVTE_CHECK message. +cudnn_frontend::error_t is_supported_f16_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, + bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, DType q_dtype, cudnnHandle_t handle); + +// Probe: same as above for the F16/BF16 backward graph. +cudnn_frontend::error_t is_supported_f16_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType q_dtype, 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 d97f388459..c9f7a9ee76 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -2991,4 +2991,105 @@ void fused_attn_fp8_bwd( return; } } + +cudnn_frontend::error_t is_supported_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, DType q_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { + // FP8 fwd impl rejects any qkv_format other than BSHD/SBHD/BHSD with NVTE_ERROR; mirror that + // here so the probe returns a typed rejection instead of catching the throw. + const 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) { + return {cudnn_frontend::error_code_t::INVALID_VALUE, + "FP8 fused attention only supports BSHD/SBHD/BHSD layouts."}; + } + size_t workspace_size = 0; + try { + fused_attn::fused_attn_fp8_fwd_impl( + static_cast(batch), static_cast(num_attn_heads), + static_cast(num_gqa_groups), static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), static_cast(head_dim_qk), + static_cast(head_dim_v), is_training, /*scaling_factor=*/1.0f, p_dropout, + qkv_layout, /*o_format=*/qkv_format, bias_type, mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, + /*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, get_cudnn_fe_dtype(q_dtype), + get_cudnn_fe_dtype(o_dtype), scaling_mode, + /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return {cudnn_frontend::error_code_t::OK, ""}; + } catch (const std::exception &e) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + } catch (...) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, + "is_supported_fp8_fwd: unknown failure"}; + } +} + +cudnn_frontend::error_t is_supported_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType q_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, cudnnHandle_t handle) { + const 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) { + return {cudnn_frontend::error_code_t::INVALID_VALUE, + "FP8 fused attention only supports BSHD/SBHD/BHSD layouts."}; + } + // For FP8 bwd, dO data type matches O data type and dQKV data type matches Q data type + // (this mirrors the assumption used by callers of fused_attn_fp8_bwd in TE). + const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(q_dtype); + const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); + const cudnn_frontend::DataType_t do_t = o_t; + const cudnn_frontend::DataType_t dqkv_t = qkv_t; + size_t workspace_size = 0; + try { + fused_attn::fused_attn_fp8_bwd_impl( + static_cast(batch), static_cast(num_attn_heads), + static_cast(num_gqa_groups), static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), static_cast(head_dim_qk), + static_cast(head_dim_v), /*scaling_factor=*/1.0f, p_dropout, qkv_layout, + /*o_format=*/qkv_format, /*do_format=*/qkv_format, /*dqkv_layout=*/qkv_layout, bias_type, + mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + deterministic, + /*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, qkv_t, o_t, do_t, dqkv_t, scaling_mode, + /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + /*do_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return {cudnn_frontend::error_code_t::OK, ""}; + } catch (const std::exception &e) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + } catch (...) { + return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, + "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 aaf5039eeb..5c7f11d80e 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,6 +8,8 @@ * \brief Functions for fused attention for FP8 with seqlen <= 512 */ +#include + #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" @@ -39,4 +41,27 @@ void fused_attn_fp8_bwd( 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); + +// Probe: drives cuDNN-FE (validate -> build_operation_graph -> create_execution_plans -> +// check_support -> build_plans) for an FP8 forward graph with the given configuration. +// Returns the cuDNN-FE status: error_code_t::OK iff the graph compiles end-to-end. On OK, +// the built graph is inserted into the same thread-local cache used by fused_attn_fp8_fwd_impl. +// On rejection, err_msg contains the underlying cuDNN-FE / NVTE_CHECK message. +cudnn_frontend::error_t is_supported_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, DType q_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + cudnnHandle_t handle); + +// Probe: same as above for the FP8 backward graph. +cudnn_frontend::error_t is_supported_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType q_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, cudnnHandle_t handle); +>>>>>>> c9006435 (refactor nvte_get_fused_attn_backend with FE calls) } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 912dc32d35..787e97d628 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" @@ -196,11 +198,40 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); +/*! \struct NVTEFusedAttnBackendStatus + * \brief Diagnostic info from \c nvte_get_fused_attn_backend. + * + * Filled by \c nvte_get_fused_attn_backend when the caller passes a non-NULL pointer. + * When the routing decision is supported, \c code is 0 and \c message is the empty + * string. When the routing rejects the configuration, \c code is the underlying + * cuDNN-FE \c cudnn_frontend::error_code_t cast to \c int (TE-synthesized post-filter + * rejections use \c INVALID_VALUE), and \c message is a null-terminated human-readable + * reason that points into per-thread storage owned by TE. The pointer is valid only + * until the next call to \c nvte_get_fused_attn_backend on the same thread. + */ +typedef struct NVTEFusedAttnBackendStatus { + int code; + const char *message; +} NVTEFusedAttnBackendStatus; + /*! \brief Get fused attention backend based on input parameters. + * + * Authoritative routing: when a non-NVTE_No_Backend value is returned, the configuration + * is guaranteed to compile through cuDNN-FE (validate -> build_operation_graph -> + * create_execution_plans -> check_support -> build_plans). The router applies a small + * set of TE-specific post-filters in addition to delegating to cuDNN-FE for capability + * checks. On success the built plan is cached, so the executor avoids rebuilding. * * \param[in] is_training Whether the model is in training mode. * \param[in] q_dtype The data type of Tensor Q. * \param[in] kv_dtype The data type of Tensors K, V. + * \param[in] o_dtype The data type of output Tensor O. Used by the FP8 + * branch to disambiguate FP8 vs HALF/BF16 output; + * ignored by the F16/BF16 branch (pass q_dtype). + * \param[in] scaling_mode Scaling mode of the input tensors. Used by the FP8 + * branch to select among delayed/current/MXFP8 recipes; + * ignored by the F16/BF16 branch + * (pass NVTE_DELAYED_TENSOR_SCALING). * \param[in] qkv_layout The layout of Tensors Q, K, V. * \param[in] bias_type The attention bias type. * \param[in] attn_mask_type The attention mask type. @@ -217,13 +248,22 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. + * \param[in] handle cuDNN handle used for the support chain. Required. + * \param[out] out_status Optional. When non-NULL, populated with a code + + * message describing why the configuration was + * rejected (NVTE_No_Backend) or with code=0 and + * message="" on success. The message buffer lives in + * thread-local storage and is overwritten on every + * call on the same thread. */ 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); + bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, + NVTEScalingMode scaling_mode, 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, cudnnHandle_t handle, + NVTEFusedAttnBackendStatus *out_status); /*! \brief Compute dot product attention with separate Q, K and V. * diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 76f2d92891..c6a8897089 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -5,6 +5,7 @@ ************************************************************************/ #include "../extensions.h" +#include "common/cudnn_utils.h" #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" @@ -17,11 +18,14 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( 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 handle = cudnnExecutionPlanManager::Instance().GetHandle(); 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); + is_training, static_cast(q_dtype), static_cast(kv_dtype), + static_cast(q_dtype), NVTE_DELAYED_TENSOR_SCALING, 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, + /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, handle, + /*out_status=*/nullptr); return backend; } @@ -272,11 +276,13 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); + auto _handle_fwd = cudnnExecutionPlanManager::Instance().GetHandle(); 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); + is_training, static_cast(dtype), static_cast(dtype), + static_cast(dtype), NVTE_DELAYED_TENSOR_SCALING, 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, /*return_max_logit=*/false, + /*cuda_graph=*/false, deterministic, _handle_fwd, /*out_status=*/nullptr); 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) */ @@ -548,11 +554,13 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); + auto _handle_bwd = cudnnExecutionPlanManager::Instance().GetHandle(); 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); + is_training, static_cast(dtype), static_cast(dtype), + static_cast(dtype), NVTE_DELAYED_TENSOR_SCALING, 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, /*return_max_logit=*/false, + /*cuda_graph=*/false, deterministic, _handle_bwd, /*out_status=*/nullptr); 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/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index e6781bd58a..d67cd4a6b9 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -6,6 +6,7 @@ #include "../extensions.h" #include "common.h" +#include "common/cudnn_utils.h" #include "pybind.h" namespace { @@ -40,17 +41,25 @@ void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &s namespace transformer_engine::pytorch { // get the fused attention backend +// +// NOTE: the underlying nvte_get_fused_attn_backend now takes o_dtype and scaling_mode in +// addition to q_dtype/kv_dtype. For the F16/BF16 routing path those are ignored, so we pass +// q_dtype as o_dtype and DELAYED_TENSOR_SCALING. This Python-facing wrapper therefore keeps +// its existing signature; FP8 callers that want authoritative routing for non-default scaling +// recipes should add o_dtype / scaling_mode parameters in a follow-up. 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) { + auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); 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); + is_training, static_cast(q_dtype), static_cast(kv_dtype), + static_cast(q_dtype), NVTE_DELAYED_TENSOR_SCALING, 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, handle, /*out_status=*/nullptr); return fused_attention_backend; } From 16b837cd0f3e27b7638bfb2a90d39056680a6b6e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 May 2026 01:58:25 +0000 Subject: [PATCH 02/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/fused_attn.cpp | 4 +-- .../fused_attn_f16_arbitrary_seqlen.cu | 34 +++++++++---------- .../common/fused_attn/fused_attn_fp8.cu | 12 +++---- .../pytorch/csrc/extensions/attention.cpp | 4 +-- 4 files changed, 26 insertions(+), 28 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 615f7c2a03..95405c0d6f 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -370,8 +370,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( auto bwd_status = is_supported_fp8_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, - o_t, scaling_mode, handle); + window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, o_t, + scaling_mode, handle); if (bwd_status.is_bad()) { set_status(out_status, bwd_status); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; 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 57ca14a3e1..1ced84755c 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 @@ -1391,11 +1391,10 @@ cudnn_frontend::error_t is_supported_f16_fwd( bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, DType q_dtype, cudnnHandle_t handle) { - const ProbeDims d = compute_probe_dims(static_cast(batch), - static_cast(num_attn_heads), - static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), qkv_layout, - bias_type); + const ProbeDims d = + compute_probe_dims(static_cast(batch), static_cast(num_attn_heads), + static_cast(max_seqlen_q), static_cast(max_seqlen_kv), + qkv_layout, bias_type); const NVTE_QKV_Format o_format = nvte_get_q_format(qkv_layout); size_t workspace_size = 0; @@ -1405,17 +1404,17 @@ cudnn_frontend::error_t is_supported_f16_fwd( static_cast(num_gqa_groups), static_cast(max_seqlen_q), static_cast(max_seqlen_kv), static_cast(head_dim_qk), static_cast(head_dim_v), d.max_b, d.max_t_q, d.max_t_kv, d.num_pages_k, - d.num_pages_v, d.page_size_k, d.page_size_v, d.max_pages_per_seq_k, - d.max_pages_per_seq_v, d.bias_b, d.bias_h, d.bias_sq, d.bias_skv, is_training, - return_max_logit, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + d.num_pages_v, d.page_size_k, d.page_size_v, d.max_pages_per_seq_k, d.max_pages_per_seq_v, + d.bias_b, d.bias_h, d.bias_sq, d.bias_skv, is_training, return_max_logit, + /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, bias_type, mask_type, + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, /*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, - get_cudnn_fe_dtype(q_dtype), /*workspace=*/nullptr, &workspace_size, + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, get_cudnn_fe_dtype(q_dtype), + /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return {cudnn_frontend::error_code_t::OK, ""}; } catch (const std::exception &e) { @@ -1432,11 +1431,10 @@ cudnn_frontend::error_t is_supported_f16_bwd( NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, DType q_dtype, cudnnHandle_t handle) { - const ProbeDims d = compute_probe_dims(static_cast(batch), - static_cast(num_attn_heads), - static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), qkv_layout, - bias_type); + const ProbeDims d = + compute_probe_dims(static_cast(batch), static_cast(num_attn_heads), + static_cast(max_seqlen_q), static_cast(max_seqlen_kv), + qkv_layout, bias_type); const NVTE_QKV_Format o_format = nvte_get_q_format(qkv_layout); const NVTE_QKV_Format do_format = o_format; const NVTE_QKV_Layout dqkv_layout = qkv_layout; @@ -1457,8 +1455,8 @@ cudnn_frontend::error_t is_supported_f16_bwd( /*devPtrdBias=*/nullptr, /*devPtrdSoftmaxOffset=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, - /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, - get_cudnn_fe_dtype(q_dtype), /*workspace=*/nullptr, &workspace_size, + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, get_cudnn_fe_dtype(q_dtype), + /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return {cudnn_frontend::error_code_t::OK, ""}; } catch (const std::exception &e) { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index c9f7a9ee76..8a152cf489 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -3014,21 +3014,21 @@ cudnn_frontend::error_t is_supported_fp8_fwd( static_cast(num_gqa_groups), static_cast(max_seqlen_q), static_cast(max_seqlen_kv), static_cast(head_dim_qk), static_cast(head_dim_v), is_training, /*scaling_factor=*/1.0f, p_dropout, - qkv_layout, /*o_format=*/qkv_format, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, + qkv_layout, /*o_format=*/qkv_format, bias_type, mask_type, softmax_type, window_size_left, + window_size_right, bottom_right_diagonal, /*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, get_cudnn_fe_dtype(q_dtype), - get_cudnn_fe_dtype(o_dtype), scaling_mode, + /*devPtrDropoutOffset=*/nullptr, get_cudnn_fe_dtype(q_dtype), get_cudnn_fe_dtype(o_dtype), + scaling_mode, /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return {cudnn_frontend::error_code_t::OK, ""}; - } catch (const std::exception &e) { + } catch (const std::exception& e) { return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; } catch (...) { return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, @@ -3084,7 +3084,7 @@ cudnn_frontend::error_t is_supported_fp8_bwd( /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return {cudnn_frontend::error_code_t::OK, ""}; - } catch (const std::exception &e) { + } catch (const std::exception& e) { return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; } catch (...) { return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index d67cd4a6b9..256ede6e55 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -58,8 +58,8 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), static_cast(q_dtype), NVTE_DELAYED_TENSOR_SCALING, 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, handle, /*out_status=*/nullptr); + max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right, return_max_logit, + cuda_graph, deterministic, handle, /*out_status=*/nullptr); return fused_attention_backend; } From 42bcd89036fc8d093a192985304018c4497b22ab Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 14:11:46 -0700 Subject: [PATCH 03/63] replace code+string with string only Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 98 +++++++++---------- .../fused_attn_f16_arbitrary_seqlen.cu | 18 ++-- .../fused_attn_f16_arbitrary_seqlen.h | 19 ++-- .../common/fused_attn/fused_attn_fp8.cu | 24 ++--- .../common/fused_attn/fused_attn_fp8.h | 18 ++-- .../include/transformer_engine/fused_attn.h | 37 +++---- .../jax/csrc/extensions/attention.cpp | 6 +- .../pytorch/csrc/extensions/attention.cpp | 2 +- 8 files changed, 103 insertions(+), 119 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 95405c0d6f..961b503c1c 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -228,28 +228,17 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { namespace { -// Per-thread storage for the message string handed back through -// NVTEFusedAttnBackendStatus::message. Re-used (cleared + re-populated) on every call to -// nvte_get_fused_attn_backend on this thread, which is exactly the lifetime documented in the -// public header. -thread_local std::string g_fused_attn_backend_status_buffer; - -// Apply (code, msg) to *out_status (if non-null), routing the message through the -// thread-local buffer so the returned `const char*` outlives this function call. -void set_status(NVTEFusedAttnBackendStatus *out_status, cudnn_frontend::error_code_t code, - const std::string &message) { - if (out_status == nullptr) return; - g_fused_attn_backend_status_buffer = message; - out_status->code = static_cast(code); - out_status->message = g_fused_attn_backend_status_buffer.c_str(); -} - -void set_status(NVTEFusedAttnBackendStatus *out_status, const cudnn_frontend::error_t &err) { - set_status(out_status, err.code, err.err_msg); -} - -void set_ok(NVTEFusedAttnBackendStatus *out_status) { - set_status(out_status, cudnn_frontend::error_code_t::OK, ""); +// Per-thread storage for the diagnostic string handed back through *out_reason. Re-used +// (cleared + re-populated) on every call to nvte_get_fused_attn_backend on this thread, +// which is exactly the lifetime documented in the public header. +thread_local std::string g_fused_attn_backend_reason_buffer; + +// Stash `reason` in the thread-local buffer and (if non-null) point *out_reason at it, +// so the returned `const char*` outlives this function call. +void set_reason(const char **out_reason, const std::string &reason) { + if (out_reason == nullptr) return; + g_fused_attn_backend_reason_buffer = reason; + *out_reason = g_fused_attn_backend_reason_buffer.c_str(); } } // namespace @@ -271,9 +260,9 @@ void set_ok(NVTEFusedAttnBackendStatus *out_status) { // executor cache-hits on. // 3. Return the selected backend, or NVTE_No_Backend if any probe rejects. // -// When `out_status` is non-null, it is filled with a code + message describing the -// rejection (or {OK, ""} on success). TE post-filter rejections synthesize an -// INVALID_VALUE entry; probe rejections forward the cuDNN-FE / NVTE_CHECK error verbatim. +// When `out_reason` is non-null, it is set to "" on success or to a tagged diagnostic +// string on rejection. TE post-filter rejections are tagged "[INVALID_VALUE] ..."; +// probe rejections forward the probe's tagged string verbatim. NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, @@ -281,11 +270,11 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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, cudnnHandle_t handle, - NVTEFusedAttnBackendStatus *out_status) { + const char **out_reason) { using namespace transformer_engine; - // Initialize to OK so callers get a clean status on the success path without us having to + // Initialize to "" so callers get a clean status on the success path without us having to // remember to set it at every return. - set_ok(out_status); + set_reason(out_reason, ""); NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); @@ -300,8 +289,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( layout_group, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v) == DType::kInt64); if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { - set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, - "Configuration requires 64-bit ragged offsets, which require cuDNN >= 9.5."); + set_reason(out_reason, + "[INVALID_VALUE] Configuration requires 64-bit ragged offsets, which require " + "cuDNN >= 9.5."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -310,8 +300,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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) { - set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, - "THD-format attention requires a padding-style mask " + set_reason(out_reason, + "[INVALID_VALUE] THD-format attention requires a padding-style mask " "(PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT)."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -324,8 +314,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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) { - set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, - "Known cuDNN <= 9.15 capture quirk: training + bshd/sbhd + " + set_reason(out_reason, + "[INVALID_VALUE] Known cuDNN <= 9.15 capture quirk: training + bshd/sbhd + " "max_seqlen_kv % 128 != 0 + cuda_graph + non-padding mask is unsupported."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -346,34 +336,34 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( if (is_fp8) { // TE-only FP8 post-filters: no 64-bit ragged offsets, no max-logit output. if (requires_64bit_ragged_offset) { - set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, - "FP8 fused attention does not support 64-bit ragged offsets."); + set_reason(out_reason, + "[INVALID_VALUE] FP8 fused attention does not support 64-bit ragged offsets."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (return_max_logit) { - set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, - "FP8 fused attention does not support return_max_logit."); + set_reason(out_reason, + "[INVALID_VALUE] FP8 fused attention does not support return_max_logit."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } const DType q_t = static_cast(q_dtype); const DType o_t = static_cast(o_dtype); - auto fwd_status = is_supported_fp8_fwd( + std::string fwd_reason = is_supported_fp8_fwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, q_t, o_t, scaling_mode, handle); - if (fwd_status.is_bad()) { - set_status(out_status, fwd_status); + if (!fwd_reason.empty()) { + set_reason(out_reason, fwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (is_training) { - auto bwd_status = is_supported_fp8_bwd( + std::string bwd_reason = is_supported_fp8_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, o_t, scaling_mode, handle); - if (bwd_status.is_bad()) { - set_status(out_status, bwd_status); + if (!bwd_reason.empty()) { + set_reason(out_reason, bwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } @@ -382,31 +372,31 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( if (is_f16_or_bf16) { const DType q_t = static_cast(q_dtype); - auto fwd_status = is_supported_f16_fwd( + std::string fwd_reason = is_supported_f16_fwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, q_t, handle); - if (fwd_status.is_bad()) { - set_status(out_status, fwd_status); + if (!fwd_reason.empty()) { + set_reason(out_reason, fwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (is_training) { - auto bwd_status = is_supported_f16_bwd( + std::string bwd_reason = is_supported_f16_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, handle); - if (bwd_status.is_bad()) { - set_status(out_status, bwd_status); + if (!bwd_reason.empty()) { + set_reason(out_reason, bwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_status(out_status, cudnn_frontend::error_code_t::INVALID_VALUE, - "Unsupported Q dtype for fused attention " + set_reason(out_reason, + "[INVALID_VALUE] Unsupported Q dtype for fused attention " "(only FP16/BF16/FP8_E4M3/FP8_E5M2 are routable)."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -500,7 +490,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso is_training, Q_type, KV_type, O_type, scaling_mode, 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, /*deterministic=*/false, handle, - /*out_status=*/nullptr); + /*out_reason=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, is_training, attn_scale, @@ -589,7 +579,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso /*is_training=*/true, Q_type, KV_type, O_type, scaling_mode, 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=*/false, cuda_graph, deterministic, - handle, /*out_status=*/nullptr); + handle, /*out_reason=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); 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 1ced84755c..70094c4e93 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 @@ -1385,7 +1385,7 @@ ProbeDims compute_probe_dims(int64_t batch, int64_t num_attn_heads, int64_t max_ } } // namespace -cudnn_frontend::error_t is_supported_f16_fwd( +std::string is_supported_f16_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, bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, @@ -1416,16 +1416,15 @@ cudnn_frontend::error_t is_supported_f16_fwd( /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, get_cudnn_fe_dtype(q_dtype), /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); - return {cudnn_frontend::error_code_t::OK, ""}; + return ""; } catch (const std::exception &e) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); } catch (...) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, - "is_supported_f16_fwd: unknown failure"}; + return "[GRAPH_NOT_SUPPORTED] is_supported_f16_fwd: unknown failure"; } } -cudnn_frontend::error_t is_supported_f16_bwd( +std::string is_supported_f16_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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, @@ -1458,12 +1457,11 @@ cudnn_frontend::error_t is_supported_f16_bwd( /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, get_cudnn_fe_dtype(q_dtype), /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); - return {cudnn_frontend::error_code_t::OK, ""}; + return ""; } catch (const std::exception &e) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); } catch (...) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, - "is_supported_f16_bwd: unknown failure"}; + return "[GRAPH_NOT_SUPPORTED] is_supported_f16_bwd: unknown failure"; } } 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 38cf48c1f0..0eabe3e8dc 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 @@ -12,7 +12,8 @@ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ #include -#include + +#include #include "common/common.h" #include "transformer_engine/fused_attn.h" @@ -50,11 +51,15 @@ void fused_attn_arbitrary_seqlen_bwd( // Probe: drives cuDNN-FE (validate -> build_operation_graph -> create_execution_plans -> // check_support -> build_plans) for an F16/BF16 forward graph with the given configuration. -// Returns the cuDNN-FE status: error_code_t::OK iff the graph compiles end-to-end. On OK, -// the built graph is inserted into the same thread-local cache used by -// fused_attn_arbitrary_seqlen_fwd_impl, so the executor cache-hits on matching descriptors. -// On rejection, err_msg contains the underlying cuDNN-FE / NVTE_CHECK message. -cudnn_frontend::error_t is_supported_f16_fwd( +// Returns an empty string iff the graph compiles end-to-end; on OK the built graph is +// inserted into the same thread-local cache used by fused_attn_arbitrary_seqlen_fwd_impl, +// so the executor cache-hits on matching descriptors. +// +// On rejection, returns a non-empty diagnostic of the form +// "[] " +// where is a stable tag mirroring cudnn_frontend::error_code_t names +// (e.g. GRAPH_NOT_SUPPORTED for cuDNN-FE rejections forwarded from the support chain). +std::string is_supported_f16_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, bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, @@ -62,7 +67,7 @@ cudnn_frontend::error_t is_supported_f16_fwd( int64_t window_size_right, bool bottom_right_diagonal, DType q_dtype, cudnnHandle_t handle); // Probe: same as above for the F16/BF16 backward graph. -cudnn_frontend::error_t is_supported_f16_bwd( +std::string is_supported_f16_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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 8a152cf489..27bd0af3f3 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -2992,7 +2992,7 @@ void fused_attn_fp8_bwd( } } -cudnn_frontend::error_t is_supported_fp8_fwd( +std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, @@ -3004,8 +3004,7 @@ cudnn_frontend::error_t is_supported_fp8_fwd( const 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) { - return {cudnn_frontend::error_code_t::INVALID_VALUE, - "FP8 fused attention only supports BSHD/SBHD/BHSD layouts."}; + return "[INVALID_VALUE] FP8 fused attention only supports BSHD/SBHD/BHSD layouts."; } size_t workspace_size = 0; try { @@ -3027,16 +3026,15 @@ cudnn_frontend::error_t is_supported_fp8_fwd( /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); - return {cudnn_frontend::error_code_t::OK, ""}; + return ""; } catch (const std::exception& e) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); } catch (...) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, - "is_supported_fp8_fwd: unknown failure"}; + return "[GRAPH_NOT_SUPPORTED] is_supported_fp8_fwd: unknown failure"; } } -cudnn_frontend::error_t is_supported_fp8_bwd( +std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, @@ -3046,8 +3044,7 @@ cudnn_frontend::error_t is_supported_fp8_bwd( const 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) { - return {cudnn_frontend::error_code_t::INVALID_VALUE, - "FP8 fused attention only supports BSHD/SBHD/BHSD layouts."}; + return "[INVALID_VALUE] FP8 fused attention only supports BSHD/SBHD/BHSD layouts."; } // For FP8 bwd, dO data type matches O data type and dQKV data type matches Q data type // (this mirrors the assumption used by callers of fused_attn_fp8_bwd in TE). @@ -3083,12 +3080,11 @@ cudnn_frontend::error_t is_supported_fp8_bwd( /*do_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); - return {cudnn_frontend::error_code_t::OK, ""}; + return ""; } catch (const std::exception& e) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, e.what()}; + return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); } catch (...) { - return {cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED, - "is_supported_fp8_bwd: unknown failure"}; + return "[GRAPH_NOT_SUPPORTED] is_supported_fp8_bwd: unknown failure"; } } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 5c7f11d80e..f91cdcf291 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,7 +8,7 @@ * \brief Functions for fused attention for FP8 with seqlen <= 512 */ -#include +#include #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" @@ -44,10 +44,15 @@ void fused_attn_fp8_bwd( // Probe: drives cuDNN-FE (validate -> build_operation_graph -> create_execution_plans -> // check_support -> build_plans) for an FP8 forward graph with the given configuration. -// Returns the cuDNN-FE status: error_code_t::OK iff the graph compiles end-to-end. On OK, -// the built graph is inserted into the same thread-local cache used by fused_attn_fp8_fwd_impl. -// On rejection, err_msg contains the underlying cuDNN-FE / NVTE_CHECK message. -cudnn_frontend::error_t is_supported_fp8_fwd( +// Returns an empty string iff the graph compiles end-to-end; on OK the built graph is +// inserted into the same thread-local cache used by fused_attn_fp8_fwd_impl. +// +// On rejection, returns a non-empty diagnostic of the form +// "[] " +// where mirrors cudnn_frontend::error_code_t names (INVALID_VALUE for the +// FP8-only layout pre-filter, GRAPH_NOT_SUPPORTED for cuDNN-FE rejections forwarded +// from the support chain). +std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, @@ -56,12 +61,11 @@ cudnn_frontend::error_t is_supported_fp8_fwd( cudnnHandle_t handle); // Probe: same as above for the FP8 backward graph. -cudnn_frontend::error_t is_supported_fp8_bwd( +std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, DType q_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); ->>>>>>> c9006435 (refactor nvte_get_fused_attn_backend with FE calls) } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 787e97d628..bbcdf08995 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -198,22 +198,6 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); -/*! \struct NVTEFusedAttnBackendStatus - * \brief Diagnostic info from \c nvte_get_fused_attn_backend. - * - * Filled by \c nvte_get_fused_attn_backend when the caller passes a non-NULL pointer. - * When the routing decision is supported, \c code is 0 and \c message is the empty - * string. When the routing rejects the configuration, \c code is the underlying - * cuDNN-FE \c cudnn_frontend::error_code_t cast to \c int (TE-synthesized post-filter - * rejections use \c INVALID_VALUE), and \c message is a null-terminated human-readable - * reason that points into per-thread storage owned by TE. The pointer is valid only - * until the next call to \c nvte_get_fused_attn_backend on the same thread. - */ -typedef struct NVTEFusedAttnBackendStatus { - int code; - const char *message; -} NVTEFusedAttnBackendStatus; - /*! \brief Get fused attention backend based on input parameters. * * Authoritative routing: when a non-NVTE_No_Backend value is returned, the configuration @@ -249,12 +233,19 @@ typedef struct NVTEFusedAttnBackendStatus { * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. * \param[in] handle cuDNN handle used for the support chain. Required. - * \param[out] out_status Optional. When non-NULL, populated with a code + - * message describing why the configuration was - * rejected (NVTE_No_Backend) or with code=0 and - * message="" on success. The message buffer lives in - * thread-local storage and is overwritten on every - * call on the same thread. + * \param[out] out_reason Optional. When non-NULL, set to a null-terminated + * diagnostic string describing why the configuration + * was rejected (NVTE_No_Backend) or set to "" on + * success. Rejection messages are tagged with a + * stable category prefix that mirrors + * \c cudnn_frontend::error_code_t, e.g. + * \c "[INVALID_VALUE] ..." for TE post-filter + * rejections and FP8 layout pre-filter rejections, + * \c "[GRAPH_NOT_SUPPORTED] ..." for cuDNN-FE + * rejections forwarded from the support chain. The + * pointer points into per-thread storage owned by TE + * and is valid only until the next call to + * \c nvte_get_fused_attn_backend on the same thread. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, @@ -263,7 +254,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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, cudnnHandle_t handle, - NVTEFusedAttnBackendStatus *out_status); + const char **out_reason); /*! \brief Compute dot product attention with separate Q, K and V. * diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index c6a8897089..669570daa5 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -25,7 +25,7 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( 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, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, handle, - /*out_status=*/nullptr); + /*out_reason=*/nullptr); return backend; } @@ -282,7 +282,7 @@ static void FusedAttnForwardImpl( static_cast(dtype), NVTE_DELAYED_TENSOR_SCALING, 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, /*return_max_logit=*/false, - /*cuda_graph=*/false, deterministic, _handle_fwd, /*out_status=*/nullptr); + /*cuda_graph=*/false, deterministic, _handle_fwd, /*out_reason=*/nullptr); 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) */ @@ -560,7 +560,7 @@ static void FusedAttnBackwardImpl( static_cast(dtype), NVTE_DELAYED_TENSOR_SCALING, 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, /*return_max_logit=*/false, - /*cuda_graph=*/false, deterministic, _handle_bwd, /*out_status=*/nullptr); + /*cuda_graph=*/false, deterministic, _handle_bwd, /*out_reason=*/nullptr); 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/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 256ede6e55..3af4ba3831 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -59,7 +59,7 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( static_cast(q_dtype), NVTE_DELAYED_TENSOR_SCALING, 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, handle, /*out_status=*/nullptr); + cuda_graph, deterministic, handle, /*out_reason=*/nullptr); return fused_attention_backend; } From de8e81457b8576b222313ae068042ff6f5b68598 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 17:05:13 -0700 Subject: [PATCH 04/63] clean up logic/comments/structure Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn.py | 4 +- .../common/fused_attn/fused_attn.cpp | 113 +++++--------- .../fused_attn_f16_arbitrary_seqlen.cu | 140 ++++++++---------- .../fused_attn_f16_arbitrary_seqlen.h | 21 +-- .../common/fused_attn/fused_attn_fp8.cu | 26 ++-- .../common/fused_attn/fused_attn_fp8.h | 21 +-- .../include/transformer_engine/fused_attn.h | 34 +---- .../common/util/pybind_helper.h | 7 + .../jax/cpp_extensions/attention.py | 20 ++- transformer_engine/jax/csrc/extensions.h | 15 +- .../jax/csrc/extensions/attention.cpp | 32 ++-- .../attention/dot_product_attention/utils.py | 9 +- transformer_engine/pytorch/csrc/extensions.h | 15 +- .../pytorch/csrc/extensions/attention.cpp | 30 ++-- 14 files changed, 209 insertions(+), 278 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 8b727b1d43..f21c6cb2d0 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -444,7 +444,7 @@ def _check_configs(self): "is either BSHD_BSHD_BSHD or THD_THD_THD" ) - self.backend = FusedAttnHelper( + self.backend, message = FusedAttnHelper( self.is_training, self.dtype, self.dtype, @@ -462,7 +462,7 @@ def _check_configs(self): (-1, -1) if self.window_size is None else self.window_size, ).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/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 961b503c1c..9587693645 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -228,41 +228,19 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { namespace { -// Per-thread storage for the diagnostic string handed back through *out_reason. Re-used -// (cleared + re-populated) on every call to nvte_get_fused_attn_backend on this thread, -// which is exactly the lifetime documented in the public header. -thread_local std::string g_fused_attn_backend_reason_buffer; - -// Stash `reason` in the thread-local buffer and (if non-null) point *out_reason at it, -// so the returned `const char*` outlives this function call. -void set_reason(const char **out_reason, const std::string &reason) { - if (out_reason == nullptr) return; - g_fused_attn_backend_reason_buffer = reason; - *out_reason = g_fused_attn_backend_reason_buffer.c_str(); +// per-thread storage for the diagnostic string +// re-used (cleared + re-populated) on every call to nvte_get_fused_attn_backend on this thread +thread_local std::string fused_attn_backend_message_buffer; + +void set_message(const char **message, const std::string &reason) { + if (message == nullptr) return; + fused_attn_backend_message_buffer = reason; + *message = fused_attn_backend_message_buffer.c_str(); } } // namespace // select a backend for fused attention -// -// Routing flow: -// 1. Apply TE post-filters that encode policies cuDNN-FE doesn't model directly: -// a. requires_64bit_ragged_offset -> cudnn >= 9.5 -// b. qkv_format == THD requires a padding-style mask -// c. cuDNN <= 9.15 + is_training + bshd/sbhd + max_seqlen_kv % 128 != 0 + -// cuda_graph + non-padding mask is rejected (known capture quirk) -// 2. Dispatch by dtype to the appropriate probe(s): -// - FP8 (E4M3/E5M2): is_supported_fp8_fwd (+ is_supported_fp8_bwd if training) -// - FP16/BF16: is_supported_f16_fwd (+ is_supported_f16_bwd if training) -// The probes call the same _impl that the executor uses, with workspace=nullptr. -// They run validate -> build_operation_graph -> create_execution_plans -> -// check_support -> build_plans, and populate a thread-local cache that the -// executor cache-hits on. -// 3. Return the selected backend, or NVTE_No_Backend if any probe rejects. -// -// When `out_reason` is non-null, it is set to "" on success or to a tagged diagnostic -// string on rejection. TE post-filter rejections are tagged "[INVALID_VALUE] ..."; -// probe rejections forward the probe's tagged string verbatim. NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, @@ -270,62 +248,50 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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, cudnnHandle_t handle, - const char **out_reason) { + const char **message) { using namespace transformer_engine; - // Initialize to "" so callers get a clean status on the success path without us having to - // remember to set it at every return. - set_reason(out_reason, ""); + set_message(message, ""); NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); const auto cudnn_runtime_version = cudnnGetVersion(); - // ---------- TE post-filters (apply before delegating to cuDNN-FE) ---------- - - // (1) Ragged-offset width: cuDNN < 9.5 only supports 32-bit offsets. + // 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); if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { - set_reason(out_reason, - "[INVALID_VALUE] Configuration requires 64-bit ragged offsets, which require " + set_message(message, + "Configuration requires 64-bit ragged offsets, which require " "cuDNN >= 9.5."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - // (2) THD requires a padding-style mask. + // THD requires padding-style mask if (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 && attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { - set_reason(out_reason, - "[INVALID_VALUE] THD-format attention requires a padding-style mask " - "(PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT)."); + set_message(message, + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - // (3) cuDNN-Graph capture quirk on cuDNN <= 9.15: training + bshd/sbhd with - // max_seqlen_kv % 128 != 0 + cuda_graph + non-padding mask hangs/miscompiles. + // avoid CUDA graph issue with cuDNN <= 9.15 if (cudnn_runtime_version <= 91500 && 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) { - set_reason(out_reason, - "[INVALID_VALUE] Known cuDNN <= 9.15 capture quirk: training + bshd/sbhd + " - "max_seqlen_kv % 128 != 0 + cuda_graph + non-padding mask is unsupported."); + set_message(message, + "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - // ---------- Dispatch by dtype ---------- - - // Probes use a single-batch graph; capability checks in cuDNN-FE are batch-agnostic. constexpr size_t probe_batch = 1; - // bottom_right_diagonal is a runtime API knob the router doesn't see; the BRCM-via-mask - // case is captured by attn_mask_type, so we probe with the default top-left alignment. constexpr bool probe_bottom_right_diagonal = false; const bool is_fp8 = @@ -334,36 +300,30 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (q_dtype == NVTEDType::kNVTEFloat16 || q_dtype == NVTEDType::kNVTEBFloat16); if (is_fp8) { - // TE-only FP8 post-filters: no 64-bit ragged offsets, no max-logit output. - if (requires_64bit_ragged_offset) { - set_reason(out_reason, - "[INVALID_VALUE] FP8 fused attention does not support 64-bit ragged offsets."); - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; - } if (return_max_logit) { - set_reason(out_reason, - "[INVALID_VALUE] FP8 fused attention does not support return_max_logit."); + set_message(message, + "FP8 fused attention does not support return_max_logit=True."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - const DType q_t = static_cast(q_dtype); + const DType qkv_t = static_cast(q_dtype); const DType o_t = static_cast(o_dtype); std::string fwd_reason = is_supported_fp8_fwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, q_t, o_t, scaling_mode, + window_size_left, window_size_right, probe_bottom_right_diagonal, qkv_t, o_t, scaling_mode, handle); if (!fwd_reason.empty()) { - set_reason(out_reason, fwd_reason); + set_message(message, fwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (is_training) { std::string bwd_reason = is_supported_fp8_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, o_t, - scaling_mode, handle); + window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, qkv_t, + o_t, scaling_mode, handle); if (!bwd_reason.empty()) { - set_reason(out_reason, bwd_reason); + set_message(message, bwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } @@ -371,33 +331,32 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } if (is_f16_or_bf16) { - const DType q_t = static_cast(q_dtype); + const DType qkv_t = static_cast(q_dtype); std::string fwd_reason = is_supported_f16_fwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, q_t, + softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, qkv_t, handle); if (!fwd_reason.empty()) { - set_reason(out_reason, fwd_reason); + set_message(message, fwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (is_training) { std::string bwd_reason = is_supported_f16_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, q_t, + window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, qkv_t, handle); if (!bwd_reason.empty()) { - set_reason(out_reason, bwd_reason); + set_message(message, bwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_reason(out_reason, - "[INVALID_VALUE] Unsupported Q dtype for fused attention " - "(only FP16/BF16/FP8_E4M3/FP8_E5M2 are routable)."); + set_message(message, + "Unsupported QKV dtype qkv_dtype=" + std::to_string(q_dtype) + " ."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -490,7 +449,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso is_training, Q_type, KV_type, O_type, scaling_mode, 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, /*deterministic=*/false, handle, - /*out_reason=*/nullptr); + /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, is_training, attn_scale, @@ -579,7 +538,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso /*is_training=*/true, Q_type, KV_type, O_type, scaling_mode, 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=*/false, cuda_graph, deterministic, - handle, /*out_reason=*/nullptr); + handle, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); 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 70094c4e93..81e66d8800 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 @@ -1334,31 +1334,18 @@ void fused_attn_arbitrary_seqlen_bwd( } } -namespace { -// Probe-time defaults for runtime-only quantities the router doesn't see (paged-KV dims, -// ragged max-tokens, bias dims). These produce a graph whose support surface matches the -// real executor's: for non-paged / non-ragged paths these are unused inside the impl; -// for ragged-THD we rebind to worst-case bounds; for paged we use 1 page of full s_kv per -// batch (= same dims as non-paged), so cuDNN-FE applies the paged-attention support rules. -struct ProbeDims { - 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; -}; - -ProbeDims compute_probe_dims(int64_t batch, int64_t num_attn_heads, int64_t max_seqlen_q, - int64_t max_seqlen_kv, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type) { +std::string is_supported_f16_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, + bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, DType qkv_dtype, + cudnnHandle_t handle) { + const auto b = static_cast(batch); + const auto h = static_cast(num_attn_heads); + const auto sq = static_cast(max_seqlen_q); + const auto skv = static_cast(max_seqlen_kv); + const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); const NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); const bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); @@ -1367,45 +1354,29 @@ ProbeDims compute_probe_dims(int64_t batch, int64_t num_attn_heads, int64_t max_ const bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); const bool has_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - ProbeDims d{}; - d.max_b = (is_ragged_q || is_ragged_kv) ? batch : 0; - d.max_t_q = is_ragged_q ? batch * max_seqlen_q : 0; - d.max_t_kv = is_ragged_kv ? batch * max_seqlen_kv : 0; - d.num_pages_k = is_paged_kv ? batch : 0; - d.num_pages_v = is_paged_kv ? batch : 0; - d.page_size_k = is_paged_kv ? max_seqlen_kv : 0; - d.page_size_v = is_paged_kv ? max_seqlen_kv : 0; - d.max_pages_per_seq_k = is_paged_kv ? 1 : 0; - d.max_pages_per_seq_v = is_paged_kv ? 1 : 0; - d.bias_b = has_bias ? batch : 0; - d.bias_h = has_bias ? num_attn_heads : 0; - d.bias_sq = has_bias ? max_seqlen_q : 0; - d.bias_skv = has_bias ? max_seqlen_kv : 0; - return d; -} -} // namespace - -std::string is_supported_f16_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, - bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, DType q_dtype, cudnnHandle_t handle) { - const ProbeDims d = - compute_probe_dims(static_cast(batch), static_cast(num_attn_heads), - static_cast(max_seqlen_q), static_cast(max_seqlen_kv), - qkv_layout, bias_type); - const NVTE_QKV_Format o_format = nvte_get_q_format(qkv_layout); + const int64_t max_b = (is_ragged_q || is_ragged_kv) ? b : 0; + const int64_t max_t_q = is_ragged_q ? b * sq : 0; + const int64_t max_t_kv = is_ragged_kv ? b * skv : 0; + const int64_t num_pages_k = is_paged_kv ? b : 0; + const int64_t num_pages_v = is_paged_kv ? b : 0; + const int64_t page_size_k = is_paged_kv ? skv : 0; + const int64_t page_size_v = is_paged_kv ? skv : 0; + const int64_t max_pages_per_seq_k = is_paged_kv ? 1 : 0; + const int64_t max_pages_per_seq_v = is_paged_kv ? 1 : 0; + const int64_t bias_b = has_bias ? b : 0; + const int64_t bias_h = has_bias ? h : 0; + const int64_t bias_sq = has_bias ? sq : 0; + const int64_t bias_skv = has_bias ? skv : 0; + + const NVTE_QKV_Format o_format = q_format; size_t workspace_size = 0; try { fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( - static_cast(batch), static_cast(num_attn_heads), - static_cast(num_gqa_groups), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(head_dim_v), d.max_b, d.max_t_q, d.max_t_kv, d.num_pages_k, - d.num_pages_v, d.page_size_k, d.page_size_v, d.max_pages_per_seq_k, d.max_pages_per_seq_v, - d.bias_b, d.bias_h, d.bias_sq, d.bias_skv, is_training, return_max_logit, + b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), + static_cast(head_dim_v), max_b, max_t_q, max_t_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, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrBias=*/nullptr, @@ -1413,14 +1384,15 @@ std::string is_supported_f16_fwd( /*devPtrO=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, /*devPtrPageTableK=*/nullptr, /*devPtrPageTableV=*/nullptr, - /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, get_cudnn_fe_dtype(q_dtype), + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, + get_cudnn_fe_dtype(qkv_dtype), /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; } catch (const std::exception &e) { - return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); + return e.what(); } catch (...) { - return "[GRAPH_NOT_SUPPORTED] is_supported_f16_fwd: unknown failure"; + return "is_supported_f16_fwd: unknown failure."; } } @@ -1429,23 +1401,36 @@ std::string is_supported_f16_bwd( size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType q_dtype, cudnnHandle_t handle) { - const ProbeDims d = - compute_probe_dims(static_cast(batch), static_cast(num_attn_heads), - static_cast(max_seqlen_q), static_cast(max_seqlen_kv), - qkv_layout, bias_type); - const NVTE_QKV_Format o_format = nvte_get_q_format(qkv_layout); + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, cudnnHandle_t handle) { + const auto b = static_cast(batch); + const auto h = static_cast(num_attn_heads); + const auto sq = static_cast(max_seqlen_q); + const auto skv = static_cast(max_seqlen_kv); + + const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(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 bool has_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + + const int64_t max_b = (is_ragged_q || is_ragged_kv) ? b : 0; + const int64_t max_t_q = is_ragged_q ? b * sq : 0; + const int64_t max_t_kv = is_ragged_kv ? b * skv : 0; + const int64_t bias_b = has_bias ? b : 0; + const int64_t bias_h = has_bias ? h : 0; + const int64_t bias_sq = has_bias ? sq : 0; + const int64_t bias_skv = has_bias ? skv : 0; + + const NVTE_QKV_Format o_format = q_format; const NVTE_QKV_Format do_format = o_format; const NVTE_QKV_Layout dqkv_layout = qkv_layout; size_t workspace_size = 0; try { fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( - static_cast(batch), static_cast(num_attn_heads), - static_cast(num_gqa_groups), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(head_dim_v), d.max_b, d.max_t_q, d.max_t_kv, d.bias_b, d.bias_h, - d.bias_sq, d.bias_skv, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, do_format, + b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), + static_cast(head_dim_v), max_b, max_t_q, max_t_kv, bias_b, bias_h, bias_sq, + bias_skv, /*scaling_factor=*/1.0f, 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=*/nullptr, /*devPtrKTranspose=*/nullptr, /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, @@ -1454,14 +1439,15 @@ std::string is_supported_f16_bwd( /*devPtrdBias=*/nullptr, /*devPtrdSoftmaxOffset=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, - /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, get_cudnn_fe_dtype(q_dtype), + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, + get_cudnn_fe_dtype(qkv_dtype), /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; } catch (const std::exception &e) { - return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); + return e.what(); } catch (...) { - return "[GRAPH_NOT_SUPPORTED] is_supported_f16_bwd: unknown failure"; + return "is_supported_f16_bwd: unknown failure."; } } 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 0eabe3e8dc..3f5ae717bb 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 @@ -49,30 +49,25 @@ void fused_attn_arbitrary_seqlen_bwd( const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); -// Probe: drives cuDNN-FE (validate -> build_operation_graph -> create_execution_plans -> -// check_support -> build_plans) for an F16/BF16 forward graph with the given configuration. -// Returns an empty string iff the graph compiles end-to-end; on OK the built graph is -// inserted into the same thread-local cache used by fused_attn_arbitrary_seqlen_fwd_impl, -// so the executor cache-hits on matching descriptors. -// -// On rejection, returns a non-empty diagnostic of the form -// "[] " -// where is a stable tag mirroring cudnn_frontend::error_code_t names -// (e.g. GRAPH_NOT_SUPPORTED for cuDNN-FE rejections forwarded from the support chain). +// 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( 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, bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, DType q_dtype, cudnnHandle_t handle); + int64_t window_size_right, bool bottom_right_diagonal, DType qkv_dtype, cudnnHandle_t handle); -// Probe: same as above for the F16/BF16 backward graph. +// 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( 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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType q_dtype, cudnnHandle_t handle); + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 27bd0af3f3..c0b515138c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -2997,14 +2997,12 @@ std::string is_supported_fp8_fwd( size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, DType q_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle) { - // FP8 fwd impl rejects any qkv_format other than BSHD/SBHD/BHSD with NVTE_ERROR; mirror that - // here so the probe returns a typed rejection instead of catching the throw. const 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) { - return "[INVALID_VALUE] FP8 fused attention only supports BSHD/SBHD/BHSD layouts."; + return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + qkv_format + "."; } size_t workspace_size = 0; try { @@ -3021,16 +3019,16 @@ std::string is_supported_fp8_fwd( /*devPtrDescaleS=*/nullptr, /*devPtrScaleS=*/nullptr, /*devPtrScaleO=*/nullptr, /*devPtrAmaxO=*/nullptr, /*devPtrAmaxS=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, - /*devPtrDropoutOffset=*/nullptr, get_cudnn_fe_dtype(q_dtype), get_cudnn_fe_dtype(o_dtype), - scaling_mode, + /*devPtrDropoutOffset=*/nullptr, get_cudnn_fe_dtype(qkv_dtype), + get_cudnn_fe_dtype(o_dtype), scaling_mode, /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; } catch (const std::exception& e) { - return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); + return e.what(); } catch (...) { - return "[GRAPH_NOT_SUPPORTED] is_supported_fp8_fwd: unknown failure"; + return "is_supported_fp8_fwd: unknown failure."; } } @@ -3039,16 +3037,14 @@ std::string is_supported_fp8_bwd( size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType q_dtype, DType o_dtype, + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle) { const 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) { - return "[INVALID_VALUE] FP8 fused attention only supports BSHD/SBHD/BHSD layouts."; + return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + qkv_format + "."; } - // For FP8 bwd, dO data type matches O data type and dQKV data type matches Q data type - // (this mirrors the assumption used by callers of fused_attn_fp8_bwd in TE). - const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(q_dtype); + const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); const cudnn_frontend::DataType_t do_t = o_t; const cudnn_frontend::DataType_t dqkv_t = qkv_t; @@ -3082,9 +3078,9 @@ std::string is_supported_fp8_bwd( /*stream=*/static_cast(0), handle); return ""; } catch (const std::exception& e) { - return std::string("[GRAPH_NOT_SUPPORTED] ") + e.what(); + return e.what(); } catch (...) { - return "[GRAPH_NOT_SUPPORTED] is_supported_fp8_bwd: unknown failure"; + return "is_supported_fp8_bwd: unknown failure."; } } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index f91cdcf291..7c7460e4ea 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -42,30 +42,25 @@ void fused_attn_fp8_bwd( 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); -// Probe: drives cuDNN-FE (validate -> build_operation_graph -> create_execution_plans -> -// check_support -> build_plans) for an FP8 forward graph with the given configuration. -// Returns an empty string iff the graph compiles end-to-end; on OK the built graph is -// inserted into the same thread-local cache used by fused_attn_fp8_fwd_impl. -// -// On rejection, returns a non-empty diagnostic of the form -// "[] " -// where mirrors cudnn_frontend::error_code_t names (INVALID_VALUE for the -// FP8-only layout pre-filter, GRAPH_NOT_SUPPORTED for cuDNN-FE rejections forwarded -// from the support chain). +// 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( 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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, DType q_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); -// Probe: same as above for the FP8 backward graph. +// 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( 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 p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType q_dtype, DType o_dtype, + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index bbcdf08995..b90749c8ee 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -199,23 +199,12 @@ 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 Get fused attention backend based on input parameters. - * - * Authoritative routing: when a non-NVTE_No_Backend value is returned, the configuration - * is guaranteed to compile through cuDNN-FE (validate -> build_operation_graph -> - * create_execution_plans -> check_support -> build_plans). The router applies a small - * set of TE-specific post-filters in addition to delegating to cuDNN-FE for capability - * checks. On success the built plan is cached, so the executor avoids rebuilding. * * \param[in] is_training Whether the model is in training mode. * \param[in] q_dtype The data type of Tensor Q. * \param[in] kv_dtype The data type of Tensors K, V. - * \param[in] o_dtype The data type of output Tensor O. Used by the FP8 - * branch to disambiguate FP8 vs HALF/BF16 output; - * ignored by the F16/BF16 branch (pass q_dtype). - * \param[in] scaling_mode Scaling mode of the input tensors. Used by the FP8 - * branch to select among delayed/current/MXFP8 recipes; - * ignored by the F16/BF16 branch - * (pass NVTE_DELAYED_TENSOR_SCALING). + * \param[in] o_dtype The data type of Tensor O. + * \param[in] scaling_mode Scaling mode of attention. * \param[in] qkv_layout The layout of Tensors Q, K, V. * \param[in] bias_type The attention bias type. * \param[in] attn_mask_type The attention mask type. @@ -232,20 +221,9 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. - * \param[in] handle cuDNN handle used for the support chain. Required. - * \param[out] out_reason Optional. When non-NULL, set to a null-terminated - * diagnostic string describing why the configuration - * was rejected (NVTE_No_Backend) or set to "" on - * success. Rejection messages are tagged with a - * stable category prefix that mirrors - * \c cudnn_frontend::error_code_t, e.g. - * \c "[INVALID_VALUE] ..." for TE post-filter - * rejections and FP8 layout pre-filter rejections, - * \c "[GRAPH_NOT_SUPPORTED] ..." for cuDNN-FE - * rejections forwarded from the support chain. The - * pointer points into per-thread storage owned by TE - * and is valid only until the next call to - * \c nvte_get_fused_attn_backend on the same thread. + * \param[in] handle cuDNN handle. + * \param[out] message Empty string on success, otherwise a diagnostic string + * describing why the configuration was rejected. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, @@ -254,7 +232,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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, cudnnHandle_t handle, - const char **out_reason); + const char **message); /*! \brief Compute dot product attention with separate Q, K and V. * diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index fdfa47da8f..fb5096b9a7 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -83,6 +83,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/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 40d02f40e1..2a38e5f6bd 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -16,7 +16,7 @@ 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 NVTE_Fused_Attn_Backend, NVTEScalingMode from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -125,14 +125,22 @@ class FusedAttnHelper: 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 + 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), + q_type, jax_dtype_to_te_dtype(self.kv_dtype), + q_type, + NVTEScalingMode.NVTE_INVALID_SCALING, self.qkv_layout.value, self.attn_bias_type.value, self.attn_mask_type.value, @@ -335,7 +343,7 @@ def abstract( out_aval = q_aval.update(shape=output_shape, dtype=q_dtype) # backend determines the softmax buffer shape/dtype - backend = FusedAttnHelper( + backend, message = FusedAttnHelper( config.is_training, q_dtype, k_dtype, @@ -372,7 +380,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 2ecfedc8a2..629b6dc3bf 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" @@ -146,12 +147,14 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, - 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); +// Returns (backend, message). `message` is empty on success, otherwise a diagnostic string +// describing why the configuration was rejected when backend = NVTE_No_Backend. +std::tuple GetFusedAttnBackend( + bool is_training, DType q_dtype, DType kv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + NVTE_QKV_Layout qkv_layout, 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); 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 669570daa5..83bddcabb1 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -12,21 +12,21 @@ namespace transformer_engine { namespace jax { -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, - 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) { +std::tuple GetFusedAttnBackend( + bool is_training, DType q_dtype, DType kv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + NVTE_QKV_Layout qkv_layout, 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 handle = cudnnExecutionPlanManager::Instance().GetHandle(); + const char *message = nullptr; auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), - static_cast(q_dtype), NVTE_DELAYED_TENSOR_SCALING, 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, - /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, handle, - /*out_reason=*/nullptr); - return backend; + static_cast(o_dtype), scaling_mode, 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, + /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, handle, &message); + return {backend, message ? std::string(message) : std::string()}; } /* @@ -279,10 +279,10 @@ static void FusedAttnForwardImpl( auto _handle_fwd = cudnnExecutionPlanManager::Instance().GetHandle(); auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), - static_cast(dtype), NVTE_DELAYED_TENSOR_SCALING, qkv_layout, bias_type, mask_type, + static_cast(dtype), NVTE_INVALID_SCALING, 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, /*return_max_logit=*/false, - /*cuda_graph=*/false, deterministic, _handle_fwd, /*out_reason=*/nullptr); + /*cuda_graph=*/false, deterministic, _handle_fwd, /*message=*/nullptr); 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) */ @@ -557,10 +557,10 @@ static void FusedAttnBackwardImpl( auto _handle_bwd = cudnnExecutionPlanManager::Instance().GetHandle(); auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), - static_cast(dtype), NVTE_DELAYED_TENSOR_SCALING, qkv_layout, bias_type, mask_type, + static_cast(dtype), NVTE_INVALID_SCALING, 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, /*return_max_logit=*/false, - /*cuda_graph=*/false, deterministic, _handle_bwd, /*out_reason=*/nullptr); + /*cuda_graph=*/false, deterministic, _handle_bwd, /*message=*/nullptr); 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/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ed87423534..f236d5a26c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1229,10 +1229,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if fp8 and fp8_meta["recipe"].fp8_dpa: q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True) kv_type = q_type - fused_attention_backend = tex.get_fused_attn_backend( + fused_attention_backend, reject_message = tex.get_fused_attn_backend( is_training, q_type, kv_type, + q_type, + tex.NVTEScalingMode.NVTE_INVALID_SCALING, QKVLayout[qkv_layout], AttnBiasType[fu_core_attention_bias_type], AttnMaskType[attn_mask_type], @@ -1251,7 +1253,10 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt deterministic, ) if fused_attention_backend == FusedAttnBackend["No_Backend"]: - logger.debug("Disabling FusedAttention as no backend supports the provided input") + logger.debug( + "Disabling FusedAttention as %s", + reject_message, + ) use_fused_attention = False fused_attention_backend = None if ( diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 4a2ea7412b..733f98e575 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -75,12 +75,15 @@ 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, const DType q_dtype, const DType kv_dtype, const DType o_dtype, + NVTEScalingMode scaling_mode, 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); 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 3af4ba3831..0c5f99ef33 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -41,26 +41,22 @@ void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &s namespace transformer_engine::pytorch { // get the fused attention backend -// -// NOTE: the underlying nvte_get_fused_attn_backend now takes o_dtype and scaling_mode in -// addition to q_dtype/kv_dtype. For the F16/BF16 routing path those are ignored, so we pass -// q_dtype as o_dtype and DELAYED_TENSOR_SCALING. This Python-facing wrapper therefore keeps -// its existing signature; FP8 callers that want authoritative routing for non-default scaling -// recipes should add o_dtype / scaling_mode parameters in a follow-up. -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) { +std::tuple get_fused_attn_backend( + bool is_training, const DType q_dtype, const DType kv_dtype, const DType o_dtype, + NVTEScalingMode scaling_mode, 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) { auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); + const char *message = nullptr; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), - static_cast(q_dtype), NVTE_DELAYED_TENSOR_SCALING, 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, handle, /*out_reason=*/nullptr); - return fused_attention_backend; + static_cast(o_dtype), scaling_mode, 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, handle, &message); + return {fused_attention_backend, message ? std::string(message) : std::string()}; } // helper function for S and dP quantizers From 81e59a9c86d8a6de3686244e7a1801e5cd3db487 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 00:11:48 +0000 Subject: [PATCH 05/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/fused_attn.cpp | 15 ++++---- .../fused_attn_f16_arbitrary_seqlen.cu | 36 ++++++++++--------- .../fused_attn_f16_arbitrary_seqlen.h | 27 +++++++------- .../common/fused_attn/fused_attn_fp8.cu | 34 +++++++++--------- .../common/fused_attn/fused_attn_fp8.h | 30 ++++++++-------- 5 files changed, 74 insertions(+), 68 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 21b0d80f4f..41607f05f7 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -263,8 +263,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( max_seqlen_kv, head_dim_qk, 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."); + "Configuration requires 64-bit ragged offsets, which require " + "cuDNN >= 9.5."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -274,7 +274,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && 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."); + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -285,8 +285,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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) { - set_message(message, - "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); + set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -300,8 +299,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( if (is_fp8) { if (return_max_logit) { - set_message(message, - "FP8 fused attention does not support return_max_logit=True."); + set_message(message, "FP8 fused attention does not support return_max_logit=True."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } const DType qkv_t = static_cast(q_dtype); @@ -354,8 +352,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_message(message, - "Unsupported QKV dtype qkv_dtype=" + std::to_string(q_dtype) + " ."); + set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(q_dtype) + " ."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } 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 81e66d8800..3a2b296ffc 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 @@ -1334,13 +1334,14 @@ void fused_attn_arbitrary_seqlen_bwd( } } -std::string is_supported_f16_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, - bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, DType qkv_dtype, - cudnnHandle_t handle) { +std::string is_supported_f16_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, bool return_max_logit, + float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + DType qkv_dtype, cudnnHandle_t handle) { const auto b = static_cast(batch); const auto h = static_cast(num_attn_heads); const auto sq = static_cast(max_seqlen_q); @@ -1375,8 +1376,8 @@ std::string is_supported_f16_fwd( fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), static_cast(head_dim_v), max_b, max_t_q, max_t_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, + 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, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrBias=*/nullptr, @@ -1396,12 +1397,13 @@ std::string is_supported_f16_fwd( } } -std::string is_supported_f16_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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, cudnnHandle_t handle) { +std::string is_supported_f16_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 p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic, DType qkv_dtype, cudnnHandle_t handle) { const auto b = static_cast(batch); const auto h = static_cast(num_attn_heads); const auto sq = static_cast(max_seqlen_q); @@ -1430,8 +1432,8 @@ std::string is_supported_f16_bwd( fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), static_cast(head_dim_v), max_b, max_t_q, max_t_kv, bias_b, bias_h, bias_sq, - bias_skv, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, do_format, - dqkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, + bias_skv, /*scaling_factor=*/1.0f, 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=*/nullptr, /*devPtrKTranspose=*/nullptr, /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, 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 3f5ae717bb..fe94d0c10c 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 @@ -52,22 +52,25 @@ void fused_attn_arbitrary_seqlen_bwd( // 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( - 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, - bool return_max_logit, float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, DType qkv_dtype, cudnnHandle_t handle); +std::string is_supported_f16_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, bool return_max_logit, + float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + DType qkv_dtype, 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( - 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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, cudnnHandle_t handle); +std::string is_supported_f16_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 p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic, DType qkv_dtype, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index db4f25c05f..40b4aa4299 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1325,13 +1325,14 @@ void fused_attn_fp8_bwd( } } -std::string is_supported_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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { +std::string is_supported_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, cudnnHandle_t handle) { const 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) { @@ -1352,8 +1353,8 @@ std::string is_supported_fp8_fwd( /*devPtrDescaleS=*/nullptr, /*devPtrScaleS=*/nullptr, /*devPtrScaleO=*/nullptr, /*devPtrAmaxO=*/nullptr, /*devPtrAmaxS=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, - /*devPtrDropoutOffset=*/nullptr, get_cudnn_fe_dtype(qkv_dtype), - get_cudnn_fe_dtype(o_dtype), scaling_mode, + /*devPtrDropoutOffset=*/nullptr, get_cudnn_fe_dtype(qkv_dtype), get_cudnn_fe_dtype(o_dtype), + scaling_mode, /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); @@ -1365,13 +1366,14 @@ std::string is_supported_fp8_fwd( } } -std::string is_supported_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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle) { +std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic, DType qkv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, cudnnHandle_t handle) { const 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) { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 96f5d54968..d52dfd246b 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -45,22 +45,24 @@ void fused_attn_fp8_bwd( // 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( - 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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle); +std::string is_supported_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 p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, 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( - 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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle); +std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic, DType qkv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, cudnnHandle_t handle); } // namespace transformer_engine From 6c5126db51cf52657eafa01d853503fd254113e2 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 17:22:09 -0700 Subject: [PATCH 06/63] fix compilation errors Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/common/fused_attn/fused_attn_fp8.cu | 6 ++++-- transformer_engine/common/fused_attn/fused_attn_fp8.h | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 40b4aa4299..842e3958bc 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1336,7 +1336,8 @@ std::string is_supported_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num const 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) { - return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + qkv_format + "."; + return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(qkv_format)) + "."; } size_t workspace_size = 0; try { @@ -1377,7 +1378,8 @@ std::string is_supported_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num const 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) { - return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + qkv_format + "."; + return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(qkv_format)) + "."; } const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index d52dfd246b..21487898a6 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -36,8 +36,8 @@ void fused_attn_fp8_bwd( 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_ZInv, - const Tensor *input_S, const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, + 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); From d35bff72911e239577e63aa99db96db7571dbb97 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 00:23:00 +0000 Subject: [PATCH 07/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/fused_attn/fused_attn_fp8.h | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 21487898a6..01c7561402 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -37,10 +37,10 @@ void fused_attn_fp8_bwd( 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 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; From f6fc58568823aae209065c387f55d59c44d2dd4f Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 18:08:56 -0700 Subject: [PATCH 08/63] remove handle from API; add bottom_right_diagonal Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn.py | 1 + .../common/fused_attn/fused_attn.cpp | 49 ++++++++++-------- .../common/fused_attn/fused_attn_fp8.cu | 10 ---- .../include/transformer_engine/fused_attn.h | 51 ++++++++++--------- transformer_engine/jax/attention.py | 12 ++++- .../jax/cpp_extensions/attention.py | 3 ++ transformer_engine/jax/csrc/extensions.h | 3 +- .../jax/csrc/extensions/attention.cpp | 19 +++---- .../attention/dot_product_attention/utils.py | 3 +- transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/attention.cpp | 8 ++- 11 files changed, 85 insertions(+), 76 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index e8da8c7366..7dc7cc4c97 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -460,6 +460,7 @@ 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(message) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 41607f05f7..9c5b91d1fc 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -246,12 +246,13 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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, cudnnHandle_t handle, + bool bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic, const char **message) { using namespace transformer_engine; set_message(message, ""); NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); const auto cudnn_runtime_version = cudnnGetVersion(); @@ -278,19 +279,10 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - // avoid CUDA graph issue with cuDNN <= 9.15 - if (cudnn_runtime_version <= 91500 && 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) { - set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; - } - + // Use batch=1 for the probe to keep graph caches minimal; batch is not part of cuDNN-FE's + // support-check criteria. All other params are passed through verbatim so the cached graph + // matches what the eventual nvte_fused_attn_fwd/bwd call will build. constexpr size_t probe_batch = 1; - constexpr bool probe_bottom_right_diagonal = false; const bool is_fp8 = (q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2); @@ -302,12 +294,18 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( set_message(message, "FP8 fused attention does not support return_max_logit=True."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } + 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; + } const DType qkv_t = static_cast(q_dtype); const DType o_t = static_cast(o_dtype); std::string fwd_reason = is_supported_fp8_fwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, qkv_t, o_t, scaling_mode, + window_size_left, window_size_right, bottom_right_diagonal, qkv_t, o_t, scaling_mode, handle); if (!fwd_reason.empty()) { set_message(message, fwd_reason); @@ -317,7 +315,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::string bwd_reason = is_supported_fp8_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, qkv_t, + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, o_t, scaling_mode, handle); if (!bwd_reason.empty()) { set_message(message, bwd_reason); @@ -328,11 +326,20 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } if (is_f16_or_bf16) { + if (cudnn_runtime_version <= 91500 && 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) { + set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } const DType qkv_t = static_cast(q_dtype); std::string fwd_reason = is_supported_f16_fwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, probe_bottom_right_diagonal, qkv_t, + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_t, handle); if (!fwd_reason.empty()) { set_message(message, fwd_reason); @@ -342,7 +349,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::string bwd_reason = is_supported_f16_bwd( probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, probe_bottom_right_diagonal, deterministic, qkv_t, + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, handle); if (!bwd_reason.empty()) { set_message(message, bwd_reason); @@ -444,8 +451,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, O_type, scaling_mode, 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, /*deterministic=*/false, handle, - /*message=*/nullptr); + window_size_right, bottom_right_diagonal, return_max_logit, cuda_graph, + /*deterministic=*/false, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { fused_attn_arbitrary_seqlen_fwd( @@ -528,8 +535,8 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( /*is_training=*/true, Q_type, KV_type, O_type, scaling_mode, 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=*/false, cuda_graph, deterministic, - handle, /*message=*/nullptr); + window_size_left, window_size_right, bottom_right_diagonal, /*return_max_logit=*/false, + cuda_graph, deterministic, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 842e3958bc..f4064a8d34 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1334,11 +1334,6 @@ std::string is_supported_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle) { const 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) { - return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + - std::to_string(static_cast(qkv_format)) + "."; - } size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_fwd_impl( @@ -1376,11 +1371,6 @@ std::string is_supported_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num bool deterministic, DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle) { const 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) { - return "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + - std::to_string(static_cast(qkv_format)) + "."; - } const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); const cudnn_frontend::DataType_t do_t = o_t; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 9bcbcc5716..85e3ea68ed 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -198,30 +198,31 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); /*! \brief Get fused attention backend based on input parameters. * - * \param[in] is_training Whether the model is in training mode. - * \param[in] q_dtype The data type of Tensor Q. - * \param[in] kv_dtype The data type of Tensors K, V. - * \param[in] o_dtype The data type of Tensor O. - * \param[in] scaling_mode Scaling mode of attention. - * \param[in] qkv_layout The layout of Tensors Q, K, V. - * \param[in] bias_type The attention bias type. - * \param[in] attn_mask_type The attention mask type. - * \param[in] softmax_type The attention softmax type. - * \param[in] dropout The dropout probability. - * \param[in] num_attn_heads The number of heads in Q. - * \param[in] num_gqa_groups The number of heads in K, V. - * \param[in] max_seqlen_q The sequence length of Q. - * \param[in] max_seqlen_kv The sequence length of K, V. - * \param[in] head_dim_qk The head dimension of Q, K. - * \param[in] head_dim_v The head dimension of V. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] return_max_logit Whether to produce Max along with Stats. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] deterministic Whether determinism is required or not. - * \param[in] handle cuDNN handle. - * \param[out] message Empty string on success, otherwise a diagnostic string - * describing why the configuration was rejected. + * \param[in] is_training Whether the model is in training mode. + * \param[in] q_dtype The data type of Tensor Q. + * \param[in] kv_dtype The data type of Tensors K, V. + * \param[in] o_dtype The data type of Tensor O. + * \param[in] scaling_mode Scaling mode of attention. + * \param[in] qkv_layout The layout of Tensors Q, K, V. + * \param[in] bias_type The attention bias type. + * \param[in] attn_mask_type The attention mask type. + * \param[in] softmax_type The attention softmax type. + * \param[in] dropout The dropout probability. + * \param[in] num_attn_heads The number of heads in Q. + * \param[in] num_gqa_groups The number of heads in K, V. + * \param[in] max_seqlen_q The sequence length of Q. + * \param[in] max_seqlen_kv The sequence length of K, V. + * \param[in] head_dim_qk The head dimension of Q, K. + * \param[in] head_dim_v The head dimension of V. + * \param[in] window_size_left Sliding window size (the left half). + * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the + * bottom right corner of the softmax matrix. + * \param[in] return_max_logit Whether to produce Max along with Stats. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. + * \param[in] deterministic Whether determinism is required or not. + * \param[out] message Empty string on success, otherwise a diagnostic string + * describing why the configuration was rejected. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, @@ -229,7 +230,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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, cudnnHandle_t handle, + bool bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic, const char **message); /*! \brief Compute dot product attention with separate Q, K and V. diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index f54a043fd2..d0e125297f 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -339,13 +339,22 @@ 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, ): """ - To check whether the fused attention kernel is supported + To check whether the fused attention kernel is supported. + + If ``bottom_right_diagonal`` is None, it is derived from the mask type, matching the + convention used everywhere else in JAX TE (see ``_FusedAttnConfig`` constructions). """ 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, q_dtype, @@ -362,6 +371,7 @@ 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() diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 1631afe4f4..e6cbb10e44 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -122,6 +122,7 @@ class FusedAttnHelper: head_dim_qk: int head_dim_v: int window_size: Tuple[int, int] + bottom_right_diagonal: bool def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -154,6 +155,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(), ) @@ -359,6 +361,7 @@ def abstract( q_head_dim, v_head_dim, config.window_size, + config.bottom_right_diagonal, ).get_fused_attn_backend() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 629b6dc3bf..d958193a7d 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -154,7 +154,8 @@ std::tuple GetFusedAttnBackend( NVTE_QKV_Layout qkv_layout, 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); + 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 4ce09368d8..d5673df8a5 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -5,7 +5,6 @@ ************************************************************************/ #include "../extensions.h" -#include "common/cudnn_utils.h" #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" @@ -17,15 +16,15 @@ std::tuple GetFusedAttnBackend( NVTE_QKV_Layout qkv_layout, 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 handle = cudnnExecutionPlanManager::Instance().GetHandle(); + size_t v_head_dim, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic) { const char *message = nullptr; auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), static_cast(o_dtype), scaling_mode, 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, - /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, handle, &message); + v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, + /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, &message); return {backend, message ? std::string(message) : std::string()}; } @@ -265,13 +264,12 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - auto _handle_fwd = cudnnExecutionPlanManager::Instance().GetHandle(); auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), static_cast(dtype), NVTE_INVALID_SCALING, 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, /*return_max_logit=*/false, - /*cuda_graph=*/false, deterministic, _handle_fwd, /*message=*/nullptr); + qk_head_dim, v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, + /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, /*message=*/nullptr); 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) */ @@ -543,13 +541,12 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); - auto _handle_bwd = cudnnExecutionPlanManager::Instance().GetHandle(); auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), static_cast(dtype), NVTE_INVALID_SCALING, 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, /*return_max_logit=*/false, - /*cuda_graph=*/false, deterministic, _handle_bwd, /*message=*/nullptr); + qk_head_dim, v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, + /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, /*message=*/nullptr); 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/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7f97a1e0f2..38542586d2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1244,13 +1244,14 @@ 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 %s", + "Disabling FusedAttention: %s", reject_message, ) use_fused_attention = False diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 733f98e575..016721f8b0 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -83,7 +83,7 @@ std::tuple get_fused_attn_backend( 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); + 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 4732d47908..2f5c7058c5 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -6,7 +6,6 @@ #include "../extensions.h" #include "common.h" -#include "common/cudnn_utils.h" #include "pybind.h" namespace { @@ -47,15 +46,14 @@ std::tuple get_fused_attn_backend( 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) { - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); + bool bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic) { const char *message = nullptr; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), static_cast(o_dtype), scaling_mode, 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, handle, &message); + head_dim_qk, head_dim_v, window_size_left, window_size_right, bottom_right_diagonal, + return_max_logit, cuda_graph, deterministic, &message); return {fused_attention_backend, message ? std::string(message) : std::string()}; } From 3e666b0c59d90a2af5ab4be38b892e1549aefd91 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 18:30:39 -0700 Subject: [PATCH 09/63] add batch_size to API Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_distributed_fused_attn.py | 3 +++ tests/jax/test_fused_attn.py | 1 + .../common/fused_attn/fused_attn.cpp | 19 +++++++------------ .../include/transformer_engine/fused_attn.h | 3 ++- transformer_engine/jax/attention.py | 2 ++ .../jax/cpp_extensions/attention.py | 4 ++++ transformer_engine/jax/csrc/extensions.h | 10 +++++----- .../jax/csrc/extensions/attention.cpp | 16 ++++++++-------- transformer_engine/jax/flax/transformer.py | 3 +++ .../attention/dot_product_attention/utils.py | 1 + transformer_engine/pytorch/csrc/extensions.h | 13 +++++++------ .../pytorch/csrc/extensions/attention.cpp | 15 ++++++++------- 12 files changed, 51 insertions(+), 39 deletions(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 50c5de1db7..39efabc598 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -75,6 +75,7 @@ def impl_test_self_attn( if not is_fused_attn_kernel_available( is_training, + batch, dtype, dtype, QKVLayout.BS3HD, @@ -227,6 +228,7 @@ def test_cross_attn( if not is_fused_attn_kernel_available( is_training, + batch, dtype, dtype, QKVLayout.BSHD_BS2HD, @@ -368,6 +370,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 7dc7cc4c97..88c485db81 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -446,6 +446,7 @@ def _check_configs(self): self.backend, message = FusedAttnHelper( self.is_training, + self.batch_size, self.dtype, self.dtype, self.qkv_layout, diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 9c5b91d1fc..e0d524c783 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -241,7 +241,7 @@ void set_message(const char **message, const std::string &reason) { // select a backend for fused attention NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, + bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, NVTEScalingMode scaling_mode, 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, @@ -279,11 +279,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - // Use batch=1 for the probe to keep graph caches minimal; batch is not part of cuDNN-FE's - // support-check criteria. All other params are passed through verbatim so the cached graph - // matches what the eventual nvte_fused_attn_fwd/bwd call will build. - constexpr size_t probe_batch = 1; - const bool is_fp8 = (q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2); const bool is_f16_or_bf16 = @@ -303,7 +298,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( const DType qkv_t = static_cast(q_dtype); const DType o_t = static_cast(o_dtype); std::string fwd_reason = is_supported_fp8_fwd( - probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_t, o_t, scaling_mode, handle); @@ -313,7 +308,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } if (is_training) { std::string bwd_reason = is_supported_fp8_bwd( - probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, o_t, scaling_mode, handle); @@ -337,7 +332,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } const DType qkv_t = static_cast(q_dtype); std::string fwd_reason = is_supported_f16_fwd( - probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_t, handle); @@ -347,7 +342,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } if (is_training) { std::string bwd_reason = is_supported_f16_bwd( - probe_batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, handle); @@ -449,7 +444,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTEScalingMode scaling_mode = input_Q->scaling_mode; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, Q_type, KV_type, O_type, scaling_mode, qkv_layout, bias_type, attn_mask_type, + is_training, b, Q_type, KV_type, O_type, scaling_mode, 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, bottom_right_diagonal, return_max_logit, cuda_graph, /*deterministic=*/false, /*message=*/nullptr); @@ -533,7 +528,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTEScalingMode scaling_mode = input_Q->scaling_mode; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - /*is_training=*/true, Q_type, KV_type, O_type, scaling_mode, qkv_layout, bias_type, + /*is_training=*/true, b, Q_type, KV_type, O_type, scaling_mode, 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, bottom_right_diagonal, /*return_max_logit=*/false, cuda_graph, deterministic, /*message=*/nullptr); diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 85e3ea68ed..227afed24e 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -199,6 +199,7 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); /*! \brief Get fused attention backend based on input parameters. * * \param[in] is_training Whether the model is in training mode. + * \param[in] batch_size Batch size. * \param[in] q_dtype The data type of Tensor Q. * \param[in] kv_dtype The data type of Tensors K, V. * \param[in] o_dtype The data type of Tensor O. @@ -225,7 +226,7 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * describing why the configuration was rejected. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, + bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, NVTEScalingMode scaling_mode, 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, diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index d0e125297f..ac6cf8975c 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -325,6 +325,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, @@ -357,6 +358,7 @@ def make_helper(attn_mask_type): ) return tex.FusedAttnHelper( is_training, + batch_size, q_dtype, kv_dtype, qkv_layout, diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index e6cbb10e44..2a533c3f3e 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -108,6 +108,7 @@ class FusedAttnHelper: """ is_training: bool + batch_size: int q_dtype: jnp.dtype kv_dtype: jnp.dtype qkv_layout: QKVLayout @@ -138,6 +139,7 @@ def get_fused_attn_backend(self): q_type = jax_dtype_to_te_dtype(self.q_dtype) return transformer_engine_jax.get_fused_attn_backend( self.is_training, + self.batch_size, q_type, jax_dtype_to_te_dtype(self.kv_dtype), q_type, @@ -345,8 +347,10 @@ def abstract( out_aval = q_aval.update(shape=output_shape, dtype=q_dtype) # backend determines the softmax buffer shape/dtype + input_batch = reduce(operator.mul, batch_shape) backend, message = FusedAttnHelper( config.is_training, + input_batch, q_dtype, k_dtype, config.qkv_layout, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index d958193a7d..1e8d99c3d8 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -150,11 +150,11 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); // Returns (backend, message). `message` is empty on success, otherwise a diagnostic string // describing why the configuration was rejected when backend = NVTE_No_Backend. std::tuple GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - NVTE_QKV_Layout qkv_layout, 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 is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, 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 bottom_right_diagonal, bool deterministic); pybind11::tuple GetFusedAttnForwardWorkspaceSizes( diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index d5673df8a5..5cd3265c3e 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -12,15 +12,15 @@ namespace transformer_engine { namespace jax { std::tuple GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - NVTE_QKV_Layout qkv_layout, 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 is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, 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 bottom_right_diagonal, bool deterministic) { const char *message = nullptr; auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(q_dtype), static_cast(kv_dtype), + is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), static_cast(o_dtype), scaling_mode, 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, bottom_right_diagonal, @@ -265,7 +265,7 @@ static void FusedAttnForwardImpl( 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), + is_training, input_batch, static_cast(dtype), static_cast(dtype), static_cast(dtype), NVTE_INVALID_SCALING, 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, bottom_right_diagonal, @@ -542,7 +542,7 @@ static void FusedAttnBackwardImpl( 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), + is_training, input_batch, static_cast(dtype), static_cast(dtype), static_cast(dtype), NVTE_INVALID_SCALING, 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, bottom_right_diagonal, diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index a2e7920843..184547aa92 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -748,6 +748,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 @@ -763,6 +765,7 @@ def __call__( has_fused_attn_kernel = 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 input_dtype, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 38542586d2..52bb687851 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1227,6 +1227,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt kv_type = q_type fused_attention_backend, reject_message = tex.get_fused_attn_backend( is_training, + batch_size, q_type, kv_type, q_type, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 016721f8b0..205e7eb834 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -78,12 +78,13 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T // 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, const DType q_dtype, const DType kv_dtype, const DType o_dtype, - NVTEScalingMode scaling_mode, 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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic); + bool is_training, size_t batch_size, const DType q_dtype, const DType kv_dtype, + const DType o_dtype, NVTEScalingMode scaling_mode, 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 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 2f5c7058c5..41dcd3301a 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -41,15 +41,16 @@ namespace transformer_engine::pytorch { // get the fused attention backend std::tuple get_fused_attn_backend( - bool is_training, const DType q_dtype, const DType kv_dtype, const DType o_dtype, - NVTEScalingMode scaling_mode, 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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic) { + bool is_training, size_t batch_size, const DType q_dtype, const DType kv_dtype, + const DType o_dtype, NVTEScalingMode scaling_mode, 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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, + bool deterministic) { const char *message = nullptr; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, static_cast(q_dtype), static_cast(kv_dtype), + is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), static_cast(o_dtype), scaling_mode, 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, bottom_right_diagonal, From 056aba6aebbfbab3e580d3155ccb07c076bcf940 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 01:31:54 +0000 Subject: [PATCH 10/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/fused_attn/fused_attn.cpp | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index e0d524c783..628bce1b54 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -310,8 +310,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::string bwd_reason = is_supported_fp8_bwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, - o_t, scaling_mode, handle); + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, o_t, + scaling_mode, handle); if (!bwd_reason.empty()) { set_message(message, bwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; @@ -334,8 +334,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::string fwd_reason = is_supported_f16_fwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_t, - handle); + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_t, handle); if (!fwd_reason.empty()) { set_message(message, fwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; @@ -344,8 +343,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::string bwd_reason = is_supported_f16_bwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, - handle); + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, handle); if (!bwd_reason.empty()) { set_message(message, bwd_reason); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; From e054863c2fda315bd43866160e188f2be95e4aea Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 20:34:12 -0700 Subject: [PATCH 11/63] fix jax binding Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/jax/csrc/extensions/pybind.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 70d0403b3e..2d55abedc6 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -206,6 +206,14 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVFP4_2D_SCALING", JAXX_Scaling_Mode::NVFP4_2D_SCALING) .export_values(); + 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, "JAXX_Quantize_Layout", pybind11::module_local()) .value("ROWWISE", JAXX_Quantize_Layout::ROWWISE) .value("COLWISE", JAXX_Quantize_Layout::COLWISE) From a7fe928eb21df1528fe57af6e87d4ff06dbd1f9b Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 22:28:35 -0700 Subject: [PATCH 12/63] specify o_dtype for FP8s Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../attention/dot_product_attention/utils.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 52bb687851..db55f2fbd3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1222,16 +1222,29 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_fused_attention: q_type = TE_DType[qkv_dtype] kv_type = q_type + o_type = q_type + scaling_mode = tex.NVTEScalingMode.NVTE_INVALID_SCALING 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 + 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] + elif recipe.float8_current_scaling() and cs_o_in_f16: + scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + o_type = TE_DType[torch.bfloat16] + else: + scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + o_type = q_type fused_attention_backend, reject_message = tex.get_fused_attn_backend( is_training, batch_size, q_type, kv_type, - q_type, - tex.NVTEScalingMode.NVTE_INVALID_SCALING, + o_type, + scaling_mode, QKVLayout[qkv_layout], AttnBiasType[fu_core_attention_bias_type], AttnMaskType[attn_mask_type], From c9b22b5d187f4cce297c54eb29976a2e53f1361a Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 8 May 2026 10:18:09 -0700 Subject: [PATCH 13/63] fix BRCM and custom_fp8 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 12 ++++++++++++ tests/pytorch/utils.py | 2 ++ 2 files changed, 14 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 32ea1694ee..4c8435f246 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2570,10 +2570,21 @@ def test_custom_mha_fp8_vs_f16(dtype, model): # 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, ) @@ -2651,6 +2662,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 3b2e50be3f..1169849044 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -275,6 +275,7 @@ 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 @@ -351,6 +352,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, From ac44e66beba56f16d1dd4c00d13aff8435025cb2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 17:19:08 +0000 Subject: [PATCH 14/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 1169849044..240acb1bd0 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -275,7 +275,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.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 From 9131b2de638ae1a40b6bae4681b2cdf73632a33e Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 8 May 2026 12:18:45 -0700 Subject: [PATCH 15/63] add o_format/etc to API and other tweaks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 65 +++++++++------- .../fused_attn_f16_arbitrary_seqlen.cu | 33 ++++---- .../fused_attn_f16_arbitrary_seqlen.h | 21 +++-- .../common/fused_attn/fused_attn_fp8.cu | 52 +++++++------ .../common/fused_attn/fused_attn_fp8.h | 30 +++++--- .../include/transformer_engine/fused_attn.h | 38 +++++++-- transformer_engine/jax/attention.py | 12 ++- .../jax/cpp_extensions/attention.py | 18 ++++- transformer_engine/jax/csrc/extensions.h | 17 ++-- .../jax/csrc/extensions/attention.cpp | 77 ++++++++++++++----- .../jax/csrc/extensions/pybind.cpp | 6 +- transformer_engine/jax/flax/transformer.py | 6 +- .../attention/dot_product_attention/utils.py | 13 ++++ transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/attention.cpp | 9 ++- 15 files changed, 274 insertions(+), 125 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 628bce1b54..47ee2ff9e5 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -231,10 +231,13 @@ namespace { // re-used (cleared + re-populated) on every call to nvte_get_fused_attn_backend on this thread thread_local std::string fused_attn_backend_message_buffer; -void set_message(const char **message, const std::string &reason) { - if (message == nullptr) return; - fused_attn_backend_message_buffer = reason; - *message = fused_attn_backend_message_buffer.c_str(); +// 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 @@ -242,12 +245,14 @@ void set_message(const char **message, const std::string &reason) { // select a backend for fused attention NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, - NVTEScalingMode scaling_mode, 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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic, - const char **message) { + 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 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, const char **message) { using namespace transformer_engine; set_message(message, ""); NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); @@ -299,21 +304,22 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( const DType o_t = static_cast(o_dtype); std::string fwd_reason = is_supported_fp8_fwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_v, is_training, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, qkv_t, o_t, scaling_mode, - handle); + head_dim_v, is_training, return_max_logit, 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, qkv_t, o_t, scaling_mode, handle); if (!fwd_reason.empty()) { - set_message(message, fwd_reason); + set_message(message, std::move(fwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (is_training) { std::string bwd_reason = is_supported_fp8_bwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, + head_dim_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, qkv_t, o_t, scaling_mode, handle); if (!bwd_reason.empty()) { - set_message(message, bwd_reason); + set_message(message, std::move(bwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } @@ -331,21 +337,25 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } const DType qkv_t = static_cast(q_dtype); + const DType o_t = static_cast(o_dtype); std::string fwd_reason = is_supported_f16_fwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_v, is_training, return_max_logit, dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_t, handle); + head_dim_v, is_training, return_max_logit, 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, qkv_t, o_t, scaling_mode, handle); if (!fwd_reason.empty()) { - set_message(message, fwd_reason); + set_message(message, std::move(fwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (is_training) { std::string bwd_reason = is_supported_f16_bwd( batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_v, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_t, handle); + head_dim_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, qkv_t, o_t, + scaling_mode, handle); if (!bwd_reason.empty()) { - set_message(message, bwd_reason); + set_message(message, std::move(bwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } @@ -442,8 +452,10 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTEScalingMode scaling_mode = input_Q->scaling_mode; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, b, Q_type, KV_type, O_type, scaling_mode, 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, + is_training, b, Q_type, KV_type, O_type, scaling_mode, qkv_layout, o_format, + /*do_format=*/o_format, /*dqkv_layout=*/qkv_layout, qkv_scale_inv_format, + /*do_scale_inv_format=*/qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, + attn_scale, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, bottom_right_diagonal, return_max_logit, cuda_graph, /*deterministic=*/false, /*message=*/nullptr); @@ -526,8 +538,9 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTEScalingMode scaling_mode = input_Q->scaling_mode; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - /*is_training=*/true, b, Q_type, KV_type, O_type, scaling_mode, qkv_layout, bias_type, - attn_mask_type, softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, + /*is_training=*/true, b, Q_type, KV_type, O_type, scaling_mode, qkv_layout, o_format, + do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, + softmax_type, attn_scale, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, bottom_right_diagonal, /*return_max_logit=*/false, cuda_graph, deterministic, /*message=*/nullptr); 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 3a2b296ffc..2b6ed2fca4 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 @@ -1337,11 +1337,15 @@ void fused_attn_arbitrary_seqlen_bwd( std::string is_supported_f16_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, bool return_max_logit, - float p_dropout, NVTE_QKV_Layout qkv_layout, + float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, + [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_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, - DType qkv_dtype, cudnnHandle_t handle) { + DType qkv_dtype, [[maybe_unused]] DType o_dtype, + [[maybe_unused]] NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { const auto b = static_cast(batch); const auto h = static_cast(num_attn_heads); const auto sq = static_cast(max_seqlen_q); @@ -1369,17 +1373,15 @@ std::string is_supported_f16_fwd(size_t batch, size_t num_attn_heads, size_t num const int64_t bias_sq = has_bias ? sq : 0; const int64_t bias_skv = has_bias ? skv : 0; - const NVTE_QKV_Format o_format = q_format; - size_t workspace_size = 0; try { fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), static_cast(head_dim_v), max_b, max_t_q, max_t_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, - /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + 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=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrS1=*/nullptr, /*devPtrS2=*/nullptr, /*devPtrO=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, @@ -1399,11 +1401,18 @@ std::string is_supported_f16_fwd(size_t batch, size_t num_attn_heads, size_t num std::string is_supported_f16_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 p_dropout, NVTE_QKV_Layout qkv_layout, + 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, + [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_format, + [[maybe_unused]] NVTE_QKV_Format do_scale_inv_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, - bool deterministic, DType qkv_dtype, cudnnHandle_t handle) { + bool deterministic, DType qkv_dtype, + [[maybe_unused]] DType o_dtype, + [[maybe_unused]] NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { const auto b = static_cast(batch); const auto h = static_cast(num_attn_heads); const auto sq = static_cast(max_seqlen_q); @@ -1423,16 +1432,12 @@ std::string is_supported_f16_bwd(size_t batch, size_t num_attn_heads, size_t num const int64_t bias_sq = has_bias ? sq : 0; const int64_t bias_skv = has_bias ? skv : 0; - const NVTE_QKV_Format o_format = q_format; - const NVTE_QKV_Format do_format = o_format; - const NVTE_QKV_Layout dqkv_layout = qkv_layout; - size_t workspace_size = 0; try { fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), static_cast(head_dim_v), max_b, max_t_q, max_t_kv, bias_b, bias_h, bias_sq, - bias_skv, /*scaling_factor=*/1.0f, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, + 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=*/nullptr, /*devPtrKTranspose=*/nullptr, /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, 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 fe94d0c10c..078b6c700d 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 @@ -55,22 +55,29 @@ void fused_attn_arbitrary_seqlen_bwd( std::string is_supported_f16_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, bool return_max_logit, - float p_dropout, NVTE_QKV_Layout qkv_layout, + 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, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - DType qkv_dtype, cudnnHandle_t handle); + DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + 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(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 p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, DType qkv_dtype, cudnnHandle_t handle); + 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, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, + DType o_dtype, NVTEScalingMode scaling_mode, + cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index f4064a8d34..fb0790f230 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1327,22 +1327,24 @@ void fused_attn_fp8_bwd( std::string is_supported_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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle) { - const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + size_t head_dim_v, bool is_training, + [[maybe_unused]] bool return_max_logit, 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, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, + DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_fwd_impl( static_cast(batch), static_cast(num_attn_heads), static_cast(num_gqa_groups), static_cast(max_seqlen_q), static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(head_dim_v), is_training, /*scaling_factor=*/1.0f, p_dropout, - qkv_layout, /*o_format=*/qkv_format, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, + static_cast(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=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, /*devPtrDescaleQ=*/nullptr, /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, @@ -1350,8 +1352,7 @@ std::string is_supported_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num /*devPtrAmaxO=*/nullptr, /*devPtrAmaxS=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, get_cudnn_fe_dtype(qkv_dtype), get_cudnn_fe_dtype(o_dtype), - scaling_mode, - /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + scaling_mode, qkv_scale_inv_format, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; @@ -1364,13 +1365,16 @@ std::string is_supported_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num std::string is_supported_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 p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle) { - const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + 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, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, + DType o_dtype, NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); const cudnn_frontend::DataType_t do_t = o_t; @@ -1381,10 +1385,9 @@ std::string is_supported_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num static_cast(batch), static_cast(num_attn_heads), static_cast(num_gqa_groups), static_cast(max_seqlen_q), static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(head_dim_v), /*scaling_factor=*/1.0f, p_dropout, qkv_layout, - /*o_format=*/qkv_format, /*do_format=*/qkv_format, /*dqkv_layout=*/qkv_layout, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, - deterministic, + static_cast(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=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, @@ -1399,8 +1402,7 @@ std::string is_supported_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num /*devPtrDescaledO_t=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, qkv_t, o_t, do_t, dqkv_t, scaling_mode, - /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, - /*do_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + qkv_scale_inv_format, do_scale_inv_format, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 01c7561402..078737b99e 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -47,22 +47,28 @@ void fused_attn_fp8_bwd( // if not, return a diagnostic message in the form of a string. std::string is_supported_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 p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle); + size_t head_dim_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_QKV_Format qkv_scale_inv_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, + DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, + 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(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 p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle); + 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, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, + DType o_dtype, NVTEScalingMode scaling_mode, + cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 227afed24e..0d20712207 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -197,6 +197,11 @@ 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 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] is_training Whether the model is in training mode. * \param[in] batch_size Batch size. @@ -205,9 +210,19 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] o_dtype The data type of Tensor O. * \param[in] scaling_mode Scaling mode of attention. * \param[in] qkv_layout The layout of Tensors Q, K, V. + * \param[in] o_format The format of Tensor O. + * \param[in] do_format The format of Tensor dO. + * \param[in] dqkv_layout The layout of Tensors dQ, dK, dV. + * \param[in] qkv_scale_inv_format Format of the scale-inverse tensors for QKV in FP8 + * configurations; pass NVTE_QKV_Format_NOT_SET to let the + * backend infer it from ``qkv_layout`` otherwise. + * \param[in] do_scale_inv_format Format of the scale-inverse tensor for dO in FP8 backward + * configurations; pass NVTE_QKV_Format_NOT_SET to let the + * backend infer it from ``do_format`` otherwise. * \param[in] bias_type The attention bias type. * \param[in] attn_mask_type The attention mask type. * \param[in] softmax_type The attention softmax type. + * \param[in] attn_scale Scaling factor for Q * K^T. * \param[in] dropout The dropout probability. * \param[in] num_attn_heads The number of heads in Q. * \param[in] num_gqa_groups The number of heads in K, V. @@ -222,17 +237,24 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. - * \param[out] message Empty string on success, otherwise a diagnostic string - * describing why the configuration was rejected. + * \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`` + * on the same thread; callers that need to retain the + * message across further calls must copy it. Pass NULL to + * skip diagnostics. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_dtype, - NVTEScalingMode scaling_mode, 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 bottom_right_diagonal, bool return_max_logit, bool cuda_graph, bool deterministic, - const char **message); + 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 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, const char **message); /*! \brief Compute dot product attention with separate Q, K and V. * diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index ac6cf8975c..735383d26d 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 @@ -341,12 +342,16 @@ def is_fused_attn_kernel_available( 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. If ``bottom_right_diagonal`` is None, it is derived from the mask type, matching the convention used everywhere else in JAX TE (see ``_FusedAttnConfig`` constructions). + + When ``return_reason`` is ``True``, returns ``(available, message)`` where ``message`` is + the diagnostic string the backend produced (empty on success). """ window_size_tuple = (-1, -1) if window_size is None else window_size @@ -376,7 +381,12 @@ def make_helper(attn_mask_type): 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 2a533c3f3e..c00feec748 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, NVTEScalingMode +from transformer_engine_jax import ( + NVTE_Fused_Attn_Backend, + NVTE_QKV_Format, + NVTE_QKV_Layout, + NVTEScalingMode, +) from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -126,7 +131,11 @@ class FusedAttnHelper: bottom_right_diagonal: bool def is_fused_attn_kernel_available(self): - """Check if there is available fused attention kernel""" + """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 @@ -145,6 +154,11 @@ def get_fused_attn_backend(self): q_type, NVTEScalingMode.NVTE_INVALID_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, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 1e8d99c3d8..813ea7db4c 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -149,13 +149,20 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); // Returns (backend, message). `message` is empty on success, otherwise a diagnostic string // describing why the configuration was rejected when backend = NVTE_No_Backend. +// `o_format`, `do_format`, and `dqkv_layout` describe the output, output-gradient, and +// QKV-gradient formats/layouts the actual fwd/bwd kernels will use; pass NVTE_QKV_Format_NOT_SET +// / NVTE_QKV_Layout_NOT_SET to request that they be inferred from `qkv_layout`. +// `qkv_scale_inv_format` / `do_scale_inv_format` describe the FP8 scale-inverse layouts; pass +// NVTE_QKV_Format_NOT_SET to let the backend infer them. std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, 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 bottom_right_diagonal, bool deterministic); + 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 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 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 5cd3265c3e..1044674856 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -13,19 +13,46 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, 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 bottom_right_diagonal, bool deterministic) { + 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 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 bottom_right_diagonal, bool deterministic) { + // For convenience, allow callers to pass *_NOT_SET sentinels and infer the missing values + // from `qkv_layout`; JAX's fused-attn path always uses matching output / dQKV layouts so + // this preserves the existing behavior without forcing every Python call site to compute them. + // The scale-inv formats stay as NOT_SET when the caller passes NOT_SET because cuDNN-frontend + // already infers them from the QKV layout for the recipes JAX currently exercises. + 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; + } + // The pointer returned via `message` aliases a thread-local buffer in libtransformer_engine that + // is overwritten by the next nvte_get_fused_attn_backend call on this thread. We copy it into a + // std::string here so the value we return is safe to retain. + // + // NOTE: attn_scale is part of the cuDNN-frontend graph cache key (FADescriptor_v1::attnScale). + // Passing 1.0f here means the graph this probe builds will not be reused at the corresponding + // FusedAttnForwardImpl/FusedAttnBackwardImpl call (which forwards the user's actual scale). + // The lost reuse is a known performance gap that will be addressed when the future + // config-struct refactor also updates this Python-facing wrapper. const char *message = nullptr; auto backend = nvte_get_fused_attn_backend( is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), - static_cast(o_dtype), scaling_mode, 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, bottom_right_diagonal, - /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, &message); - return {backend, message ? std::string(message) : std::string()}; + static_cast(o_dtype), scaling_mode, qkv_layout, o_format, do_format, dqkv_layout, + qkv_scale_inv_format, do_scale_inv_format, bias_type, mask_type, softmax_type, + /*attn_scale=*/1.0f, 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, + bottom_right_diagonal, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, + &message); + return {backend, message != nullptr ? std::string(message) : std::string()}; } /* @@ -264,12 +291,19 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); + // JAX uses the same layout for output / dQKV as for QKV, so derive the formats from qkv_layout. + // Scale-inv formats stay NOT_SET because JAX's fused-attn path here is non-FP8. + const NVTE_QKV_Format probe_o_format = nvte_get_q_format(qkv_layout); auto backend = nvte_get_fused_attn_backend( is_training, input_batch, static_cast(dtype), static_cast(dtype), - static_cast(dtype), NVTE_INVALID_SCALING, 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, bottom_right_diagonal, - /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, /*message=*/nullptr); + static_cast(dtype), NVTE_INVALID_SCALING, qkv_layout, probe_o_format, + /*do_format=*/probe_o_format, /*dqkv_layout=*/qkv_layout, + /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + /*do_scale_inv_format=*/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, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, + /*message=*/nullptr); 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) */ @@ -541,12 +575,19 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); + // JAX uses the same layout for output / dQKV as for QKV, so derive the formats from qkv_layout. + // Scale-inv formats stay NOT_SET because JAX's fused-attn path here is non-FP8. + const NVTE_QKV_Format probe_o_format = nvte_get_q_format(qkv_layout); auto backend = nvte_get_fused_attn_backend( is_training, input_batch, static_cast(dtype), static_cast(dtype), - static_cast(dtype), NVTE_INVALID_SCALING, 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, bottom_right_diagonal, - /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, /*message=*/nullptr); + static_cast(dtype), NVTE_INVALID_SCALING, qkv_layout, probe_o_format, + /*do_format=*/probe_o_format, /*dqkv_layout=*/qkv_layout, + /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + /*do_scale_inv_format=*/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, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, + /*message=*/nullptr); 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 2d55abedc6..bdfec12b8b 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -160,12 +160,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 184547aa92..3b8682d7bc 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -762,7 +762,7 @@ def __call__( head_dim_qk = self.head_dim head_dim_v = self.head_dim - 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, @@ -781,15 +781,17 @@ def __call__( head_dim_qk, head_dim_v, self.window_size, + return_reason=True, ) 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"Reason for this rejection is: {reason}\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" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index db55f2fbd3..cf2f297083 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, @@ -1224,6 +1225,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt kv_type = q_type o_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: recipe = fp8_meta["recipe"] q_type = get_fp8_te_dtype(recipe, fprop_tensor=True) @@ -1232,12 +1235,17 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if recipe.mxfp8(): scaling_mode = tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING o_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] else: scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING o_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, @@ -1246,6 +1254,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt o_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], diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 205e7eb834..74021b81b5 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -80,6 +80,8 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T 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, 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 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, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 41dcd3301a..4cda724a8b 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -43,6 +43,8 @@ namespace transformer_engine::pytorch { 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, 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 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, @@ -51,11 +53,12 @@ std::tuple get_fused_attn_backend( const char *message = nullptr; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), - static_cast(o_dtype), scaling_mode, qkv_layout, bias_type, attn_mask_type, - softmax_type, p_dropout, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, + static_cast(o_dtype), scaling_mode, qkv_layout, o_format, do_format, dqkv_layout, + qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, softmax_type, + /*attn_scale=*/1.0f, 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, bottom_right_diagonal, return_max_logit, cuda_graph, deterministic, &message); - return {fused_attention_backend, message ? std::string(message) : std::string()}; + return {fused_attention_backend, message != nullptr ? std::string(message) : std::string()}; } // helper function for S and dP quantizers From 956f159d794ae11e53a8eee5092a80a2dbf3a525 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 19:20:12 +0000 Subject: [PATCH 16/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/fused_attn.cpp | 14 ++--- .../fused_attn_f16_arbitrary_seqlen.cu | 51 ++++++++----------- .../fused_attn_f16_arbitrary_seqlen.h | 3 +- .../common/fused_attn/fused_attn_fp8.cu | 46 ++++++++--------- .../common/fused_attn/fused_attn_fp8.h | 3 +- .../include/transformer_engine/fused_attn.h | 14 ++--- transformer_engine/jax/csrc/extensions.h | 12 ++--- .../jax/csrc/extensions/attention.cpp | 12 ++--- 8 files changed, 70 insertions(+), 85 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 47ee2ff9e5..4c64d9a3d8 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -246,13 +246,13 @@ void set_message(const char **message, std::string reason) { NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_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 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, const char **message) { + 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 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, + const char **message) { using namespace transformer_engine; set_message(message, ""); NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); 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 2b6ed2fca4..372bc68288 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 @@ -1334,18 +1334,15 @@ void fused_attn_arbitrary_seqlen_bwd( } } -std::string is_supported_f16_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, bool return_max_logit, - float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, - [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_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, - DType qkv_dtype, [[maybe_unused]] DType o_dtype, - [[maybe_unused]] NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { +std::string is_supported_f16_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, + bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_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, + DType qkv_dtype, [[maybe_unused]] DType o_dtype, [[maybe_unused]] NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { const auto b = static_cast(batch); const auto h = static_cast(num_attn_heads); const auto sq = static_cast(max_seqlen_q); @@ -1399,20 +1396,16 @@ std::string is_supported_f16_fwd(size_t batch, size_t num_attn_heads, size_t num } } -std::string is_supported_f16_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, - [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_format, - [[maybe_unused]] NVTE_QKV_Format do_scale_inv_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, - bool deterministic, DType qkv_dtype, - [[maybe_unused]] DType o_dtype, - [[maybe_unused]] NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { +std::string is_supported_f16_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, [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_format, + [[maybe_unused]] NVTE_QKV_Format do_scale_inv_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, bool deterministic, DType qkv_dtype, + [[maybe_unused]] DType o_dtype, [[maybe_unused]] NVTEScalingMode scaling_mode, + cudnnHandle_t handle) { const auto b = static_cast(batch); const auto h = static_cast(num_attn_heads); const auto sq = static_cast(max_seqlen_q); @@ -1437,9 +1430,9 @@ std::string is_supported_f16_bwd(size_t batch, size_t num_attn_heads, size_t num fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), static_cast(head_dim_v), max_b, max_t_q, max_t_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=*/nullptr, /*devPtrKTranspose=*/nullptr, + 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=*/nullptr, /*devPtrKTranspose=*/nullptr, /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, /*devPtrdO=*/nullptr, 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 078b6c700d..9d6b57d0a0 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 @@ -76,8 +76,7 @@ std::string is_supported_f16_bwd(size_t batch, size_t num_attn_heads, size_t num 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, DType qkv_dtype, - DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle); + DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index fb0790f230..352e47fbfd 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1325,17 +1325,14 @@ void fused_attn_fp8_bwd( } } -std::string is_supported_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, - [[maybe_unused]] bool return_max_logit, 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, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, - DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { +std::string is_supported_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, + [[maybe_unused]] bool return_max_logit, 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, + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, + DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle) { size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_fwd_impl( @@ -1363,18 +1360,15 @@ std::string is_supported_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num } } -std::string is_supported_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, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, - DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { +std::string is_supported_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, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, DType o_dtype, + NVTEScalingMode scaling_mode, cudnnHandle_t handle) { const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); const cudnn_frontend::DataType_t do_t = o_t; @@ -1385,9 +1379,9 @@ std::string is_supported_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num static_cast(batch), static_cast(num_attn_heads), static_cast(num_gqa_groups), static_cast(max_seqlen_q), static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(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, + static_cast(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=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 078737b99e..8dfdb8c412 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -69,6 +69,5 @@ std::string is_supported_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num 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, DType qkv_dtype, - DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle); + DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 0d20712207..8e9864f916 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -248,13 +248,13 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_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 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, const char **message); + 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 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, + const char **message); /*! \brief Compute dot product attention with separate Q, K and V. * diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 813ea7db4c..c543ae6019 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -157,12 +157,12 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 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 bottom_right_diagonal, bool deterministic); + 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 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 1044674856..498d614812 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -14,12 +14,12 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 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 bottom_right_diagonal, bool deterministic) { + 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 bottom_right_diagonal, bool deterministic) { // For convenience, allow callers to pass *_NOT_SET sentinels and infer the missing values // from `qkv_layout`; JAX's fused-attn path always uses matching output / dQKV layouts so // this preserves the existing behavior without forcing every Python call site to compute them. From b21f6065a1cd900b36cfd2cc8879074718646c76 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 8 May 2026 12:33:51 -0700 Subject: [PATCH 17/63] minor tweaks for docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/jax/attention.py | 5 +--- transformer_engine/jax/csrc/extensions.h | 7 ----- .../jax/csrc/extensions/attention.cpp | 28 ++++--------------- transformer_engine/jax/flax/transformer.py | 5 ++-- 4 files changed, 9 insertions(+), 36 deletions(-) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 735383d26d..e4fce42ce7 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -347,11 +347,8 @@ def is_fused_attn_kernel_available( """ To check whether the fused attention kernel is supported. - If ``bottom_right_diagonal`` is None, it is derived from the mask type, matching the - convention used everywhere else in JAX TE (see ``_FusedAttnConfig`` constructions). - When ``return_reason`` is ``True``, returns ``(available, message)`` where ``message`` is - the diagnostic string the backend produced (empty on success). + 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 diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index c543ae6019..cf455357c6 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -147,13 +147,6 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); -// Returns (backend, message). `message` is empty on success, otherwise a diagnostic string -// describing why the configuration was rejected when backend = NVTE_No_Backend. -// `o_format`, `do_format`, and `dqkv_layout` describe the output, output-gradient, and -// QKV-gradient formats/layouts the actual fwd/bwd kernels will use; pass NVTE_QKV_Format_NOT_SET -// / NVTE_QKV_Layout_NOT_SET to request that they be inferred from `qkv_layout`. -// `qkv_scale_inv_format` / `do_scale_inv_format` describe the FP8 scale-inverse layouts; pass -// NVTE_QKV_Format_NOT_SET to let the backend infer them. std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 498d614812..e37bd4442c 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -14,17 +14,12 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 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 bottom_right_diagonal, bool deterministic) { - // For convenience, allow callers to pass *_NOT_SET sentinels and infer the missing values - // from `qkv_layout`; JAX's fused-attn path always uses matching output / dQKV layouts so - // this preserves the existing behavior without forcing every Python call site to compute them. - // The scale-inv formats stay as NOT_SET when the caller passes NOT_SET because cuDNN-frontend - // already infers them from the QKV layout for the recipes JAX currently exercises. + 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 bottom_right_diagonal, bool deterministic) { if (o_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { o_format = nvte_get_q_format(qkv_layout); } @@ -34,15 +29,6 @@ std::tuple GetFusedAttnBackend( if (dqkv_layout == NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET) { dqkv_layout = qkv_layout; } - // The pointer returned via `message` aliases a thread-local buffer in libtransformer_engine that - // is overwritten by the next nvte_get_fused_attn_backend call on this thread. We copy it into a - // std::string here so the value we return is safe to retain. - // - // NOTE: attn_scale is part of the cuDNN-frontend graph cache key (FADescriptor_v1::attnScale). - // Passing 1.0f here means the graph this probe builds will not be reused at the corresponding - // FusedAttnForwardImpl/FusedAttnBackwardImpl call (which forwards the user's actual scale). - // The lost reuse is a known performance gap that will be addressed when the future - // config-struct refactor also updates this Python-facing wrapper. const char *message = nullptr; auto backend = nvte_get_fused_attn_backend( is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), @@ -291,8 +277,6 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - // JAX uses the same layout for output / dQKV as for QKV, so derive the formats from qkv_layout. - // Scale-inv formats stay NOT_SET because JAX's fused-attn path here is non-FP8. const NVTE_QKV_Format probe_o_format = nvte_get_q_format(qkv_layout); auto backend = nvte_get_fused_attn_backend( is_training, input_batch, static_cast(dtype), static_cast(dtype), diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 3b8682d7bc..f5ca3ff04d 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -789,12 +789,11 @@ def __call__( 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" - f"Reason for this rejection is: {reason}\n" + "Falling back to the unfused attention backend as fused attention does not support:\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" + f"Reason for this rejection: {reason}\n" ) dropout_rng = None From 34219205eab29cfc68c0afcda53b44309edd53c6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 19:36:47 +0000 Subject: [PATCH 18/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/jax/csrc/extensions/attention.cpp | 12 ++++++------ transformer_engine/jax/flax/transformer.py | 8 +++----- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index e37bd4442c..9c6f2483cd 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -14,12 +14,12 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 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 bottom_right_diagonal, bool deterministic) { + 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 bottom_right_diagonal, bool deterministic) { if (o_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { o_format = nvte_get_q_format(qkv_layout); } diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index f5ca3ff04d..35a48442d2 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -789,11 +789,9 @@ def __call__( if enable_fused_attn and not has_fused_attn_kernel: reason = fused_attn_reject_reason or "(no diagnostic message available)" warnings.warn( - "Falling back to the unfused attention backend as fused attention does not support:\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" - f"Reason for this rejection: {reason}\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 From 7956b4339ccf44f4d4792dd02a378c3520ef7b8a Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 8 May 2026 14:41:34 -0700 Subject: [PATCH 19/63] replace with nvte_get_fused_attn_backend_v2 and add NVTEFusedAttnConfig Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 231 ++++++++++++------ .../fused_attn_f16_arbitrary_seqlen.cu | 70 ++++-- .../fused_attn_f16_arbitrary_seqlen.h | 23 +- .../common/fused_attn/fused_attn_fp8.cu | 68 ++++-- .../common/fused_attn/fused_attn_fp8.h | 23 +- .../include/transformer_engine/fused_attn.h | 185 +++++++++----- .../jax/cpp_extensions/attention.py | 3 + transformer_engine/jax/csrc/extensions.h | 6 +- .../jax/csrc/extensions/attention.cpp | 92 ++++--- .../dot_product_attention.py | 2 + .../attention/dot_product_attention/utils.py | 6 + transformer_engine/pytorch/csrc/extensions.h | 8 +- .../pytorch/csrc/extensions/attention.cpp | 48 +++- 13 files changed, 496 insertions(+), 269 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 4c64d9a3d8..6670fd59ed 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -228,7 +228,7 @@ 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 on this thread +// 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, @@ -243,30 +243,26 @@ void set_message(const char **message, std::string reason) { } // namespace // select a backend for fused attention -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_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 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, - const char **message) { +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig *cfg, + const char **message) { using namespace transformer_engine; set_message(message, ""); - NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + NVTE_CHECK(cfg != nullptr, "NVTEFusedAttnConfig pointer must not be NULL."); + NVTE_CHECK(cfg->struct_size == sizeof(NVTEFusedAttnConfig), + "NVTEFusedAttnConfig::struct_size must equal sizeof(NVTEFusedAttnConfig); " + "did you forget NVTE_FUSED_ATTN_CONFIG_INIT?"); cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + 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(); // 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); + (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 " @@ -276,21 +272,21 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // THD requires padding-style mask if (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 && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { + 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 = - (q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2); - const bool is_f16_or_bf16 = - (q_dtype == NVTEDType::kNVTEFloat16 || q_dtype == NVTEDType::kNVTEBFloat16); + 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 (return_max_logit) { + 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; } @@ -300,24 +296,13 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::to_string(static_cast(qkv_format)) + "."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - const DType qkv_t = static_cast(q_dtype); - const DType o_t = static_cast(o_dtype); - std::string fwd_reason = is_supported_fp8_fwd( - batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_v, is_training, return_max_logit, 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, qkv_t, o_t, scaling_mode, handle); + 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 (is_training) { - std::string bwd_reason = is_supported_fp8_bwd( - batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_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, qkv_t, o_t, - scaling_mode, handle); + 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; @@ -327,33 +312,22 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } if (is_f16_or_bf16) { - if (cudnn_runtime_version <= 91500 && is_training && + 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) { + (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; } - const DType qkv_t = static_cast(q_dtype); - const DType o_t = static_cast(o_dtype); - std::string fwd_reason = is_supported_f16_fwd( - batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_v, is_training, return_max_logit, 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, qkv_t, o_t, scaling_mode, handle); + 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 (is_training) { - std::string bwd_reason = is_supported_f16_bwd( - batch_size, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, - head_dim_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, qkv_t, o_t, - scaling_mode, handle); + 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; @@ -362,10 +336,48 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(q_dtype) + " ."); + set_message(message, + "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg->qkv_dtype) + " ."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } +// 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, 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) { + NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + cfg.qkv_layout = qkv_layout; + cfg.dqkv_layout = qkv_layout; // legacy: gradient layout matches input 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 = is_training; + cfg.return_max_logit = return_max_logit; + cfg.deterministic = deterministic; + return nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); +} + // NVTE fused attention FWD with separate Q, K and V void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, @@ -448,16 +460,61 @@ 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; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, b, Q_type, KV_type, O_type, scaling_mode, qkv_layout, o_format, - /*do_format=*/o_format, /*dqkv_layout=*/qkv_layout, qkv_scale_inv_format, - /*do_scale_inv_format=*/qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, - attn_scale, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, - window_size_right, bottom_right_diagonal, return_max_logit, cuda_graph, - /*deterministic=*/false, /*message=*/nullptr); + size_t bias_b = 0, bias_h = 0, bias_sq = 0, bias_skv = 0; + if (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]; + } + + NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + cfg.qkv_layout = qkv_layout; + cfg.o_format = o_format; + cfg.do_format = o_format; // fwd path: same format used for dO if/when probed for bwd + cfg.dqkv_layout = qkv_layout; // fwd path: same layout used for dQKV if/when probed for bwd + cfg.qkv_scale_inv_format = qkv_scale_inv_format; + cfg.do_scale_inv_format = qkv_scale_inv_format; // fwd path: mirror QKV + 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.cuda_graph = cuda_graph; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = O_type; + cfg.do_dtype = O_type; // fwd path: dO assumed to share dtype with O + cfg.dqkv_dtype = Q_type; // fwd path: dQKV assumed to share dtype with QKV + 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.is_training = is_training; + cfg.return_max_logit = return_max_logit; + cfg.deterministic = false; + NVTE_Fused_Attn_Backend fused_attention_backend = + nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { fused_attn_arbitrary_seqlen_fwd( @@ -534,15 +591,45 @@ 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; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - /*is_training=*/true, b, Q_type, KV_type, O_type, scaling_mode, qkv_layout, o_format, - do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, - softmax_type, attn_scale, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, - window_size_left, window_size_right, bottom_right_diagonal, /*return_max_logit=*/false, - cuda_graph, deterministic, /*message=*/nullptr); + NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + 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.cuda_graph = cuda_graph; + 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.is_training = true; + cfg.return_max_logit = false; + cfg.deterministic = deterministic; + NVTE_Fused_Attn_Backend fused_attention_backend = + nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; 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 372bc68288..9cdee256ed 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 @@ -1334,19 +1334,27 @@ void fused_attn_arbitrary_seqlen_bwd( } } -std::string is_supported_f16_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, - bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_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, - DType qkv_dtype, [[maybe_unused]] DType o_dtype, [[maybe_unused]] NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { - const auto b = static_cast(batch); - const auto h = static_cast(num_attn_heads); - const auto sq = static_cast(max_seqlen_q); - const auto skv = static_cast(max_seqlen_kv); +std::string is_supported_f16_fwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { + const size_t num_gqa_groups = cfg->num_gqa_groups; + const size_t head_dim_qk = cfg->head_dim_qk; + const size_t head_dim_v = cfg->head_dim_v; + const bool is_training = cfg->is_training; + const bool return_max_logit = cfg->return_max_logit; + const float attn_scale = cfg->attn_scale; + const float p_dropout = 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 DType qkv_dtype = static_cast(cfg->qkv_dtype); + const auto b = static_cast(cfg->batch_size); + const auto h = static_cast(cfg->num_attn_heads); + const auto sq = static_cast(cfg->max_seqlen_q); + const auto skv = static_cast(cfg->max_seqlen_kv); const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); const NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); @@ -1396,20 +1404,28 @@ std::string is_supported_f16_fwd( } } -std::string is_supported_f16_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, [[maybe_unused]] NVTE_QKV_Format qkv_scale_inv_format, - [[maybe_unused]] NVTE_QKV_Format do_scale_inv_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, bool deterministic, DType qkv_dtype, - [[maybe_unused]] DType o_dtype, [[maybe_unused]] NVTEScalingMode scaling_mode, - cudnnHandle_t handle) { - const auto b = static_cast(batch); - const auto h = static_cast(num_attn_heads); - const auto sq = static_cast(max_seqlen_q); - const auto skv = static_cast(max_seqlen_kv); +std::string is_supported_f16_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { + const size_t num_gqa_groups = cfg->num_gqa_groups; + const size_t head_dim_qk = cfg->head_dim_qk; + const size_t head_dim_v = cfg->head_dim_v; + const float attn_scale = cfg->attn_scale; + const float p_dropout = 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 DType qkv_dtype = static_cast(cfg->qkv_dtype); + const auto b = static_cast(cfg->batch_size); + const auto h = static_cast(cfg->num_attn_heads); + const auto sq = static_cast(cfg->max_seqlen_q); + const auto skv = static_cast(cfg->max_seqlen_kv); const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); const NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); 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 9d6b57d0a0..5d27e82278 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 @@ -52,31 +52,12 @@ void fused_attn_arbitrary_seqlen_bwd( // 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(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, bool return_max_logit, - 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, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, - DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle); +std::string is_supported_f16_fwd(const NVTEFusedAttnConfig *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(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, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, - DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); +std::string is_supported_f16_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 352e47fbfd..3fe5b7fb10 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1325,14 +1325,30 @@ void fused_attn_fp8_bwd( } } -std::string is_supported_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, - [[maybe_unused]] bool return_max_logit, 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, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle) { +std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { + 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 bool is_training = cfg->is_training; + const float attn_scale = cfg->attn_scale; + const float p_dropout = cfg->dropout; + const NVTE_QKV_Layout qkv_layout = cfg->qkv_layout; + const NVTE_QKV_Format o_format = cfg->o_format; + const NVTE_QKV_Format qkv_scale_inv_format = cfg->qkv_scale_inv_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 DType qkv_dtype = static_cast(cfg->qkv_dtype); + const DType o_dtype = static_cast(cfg->o_dtype); + const NVTEScalingMode scaling_mode = cfg->scaling_mode; + size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_fwd_impl( @@ -1360,15 +1376,33 @@ std::string is_supported_fp8_fwd( } } -std::string is_supported_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, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, DType o_dtype, - NVTEScalingMode scaling_mode, cudnnHandle_t handle) { +std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { + 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 float attn_scale = cfg->attn_scale; + const float p_dropout = 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_QKV_Format qkv_scale_inv_format = cfg->qkv_scale_inv_format; + const NVTE_QKV_Format do_scale_inv_format = cfg->do_scale_inv_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 bool deterministic = cfg->deterministic; + const DType qkv_dtype = static_cast(cfg->qkv_dtype); + const DType o_dtype = static_cast(cfg->o_dtype); + const NVTEScalingMode scaling_mode = cfg->scaling_mode; + const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); const cudnn_frontend::DataType_t do_t = o_t; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 8dfdb8c412..fc60987cf3 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -45,29 +45,10 @@ void fused_attn_fp8_bwd( // 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(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, bool return_max_logit, - 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, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, - DType qkv_dtype, DType o_dtype, NVTEScalingMode scaling_mode, - cudnnHandle_t handle); +std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig *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(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, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, DType qkv_dtype, - DType o_dtype, NVTEScalingMode scaling_mode, cudnnHandle_t handle); +std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 8e9864f916..df15148350 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -196,65 +196,140 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); +/*! \struct NVTEFusedAttnConfig + * \brief Attention configuration. + * + * Versioning rules: + * - ``struct_size`` MUST be set to ``sizeof(NVTEFusedAttnConfig)`` by the + * caller (use ``NVTE_FUSED_ATTN_CONFIG_INIT``). + * - New fields may only be appended at the end; existing fields are never + * reordered, removed, or resized. The library reads only fields that are + * in range according to ``struct_size`` and uses safe defaults otherwise. + */ +typedef struct NVTEFusedAttnConfig { + size_t struct_size; /*!< MUST equal sizeof(NVTEFusedAttnConfig). */ + uint32_t reserved0; /*!< Padding for layout stability; set to 0. */ + uint32_t reserved1; /*!< Padding for layout stability; set to 0. */ + + NVTE_QKV_Layout qkv_layout; /*!< QKV tensors' layout. */ + NVTE_QKV_Format o_format; /*!< Output O tensor format. */ + NVTE_QKV_Format do_format; /*!< Output-grad dO tensor format (bwd). */ + NVTE_QKV_Layout dqkv_layout; /*!< Gradient dQKV tensor layout (bwd). */ + NVTE_QKV_Format qkv_scale_inv_format; /*!< QKV scale_inv tensor format (FP8). */ + NVTE_QKV_Format do_scale_inv_format; /*!< dO scale_inv tensor format (FP8 bwd). */ + NVTE_Bias_Type bias_type; /*!< Attention bias type. */ + NVTE_Mask_Type attn_mask_type; /*!< Attention mask type. */ + NVTE_Softmax_Type softmax_type; /*!< Attention softmax type. */ + NVTEScalingMode scaling_mode; /*!< Scaling mode (e.g. delayed, MXFP8). */ + float attn_scale; /*!< Pre-softmax attention scale factor. */ + float dropout; /*!< Dropout probability. */ + size_t max_seqlen_q; /*!< Max sequence length for Q. */ + size_t max_seqlen_kv; /*!< Max sequence length for K, V. */ + int64_t window_size_left; /*!< Sliding window size (left half); -1 = unlimited. */ + int64_t window_size_right; /*!< Sliding window size (right half); -1 = unlimited. */ + bool bottom_right_diagonal; /*!< Whether causal mask aligns to the bottom-right diagonal. */ + bool cuda_graph; /*!< Whether CUDA graph capture is enabled. */ + + NVTEDType qkv_dtype; /*!< Data type of Tensors Q, K, V. Q and K/V must share a dtype. */ + NVTEDType o_dtype; /*!< Data type of Tensor O. */ + NVTEDType do_dtype; /*!< Data type of Tensor dO (bwd). */ + NVTEDType dqkv_dtype; /*!< Data type of Tensors dQ, dK, dV (bwd). */ + size_t batch_size; /*!< Batch size. */ + size_t num_attn_heads; /*!< Number of heads in Q. */ + size_t num_gqa_groups; /*!< Number of heads in K, V. */ + size_t head_dim_qk; /*!< Head dimension of Q, K. */ + size_t head_dim_v; /*!< Head dimension of V. */ + + size_t num_pages_k; /*!< Total number of K cache pages. */ + size_t num_pages_v; /*!< Total number of V cache pages. */ + size_t page_size_k; /*!< Tokens per K cache page. */ + size_t page_size_v; /*!< Tokens per V cache page. */ + size_t max_pages_per_seq_k; /*!< Max K pages per sequence in the batch. */ + size_t max_pages_per_seq_v; /*!< Max V pages per sequence in the batch. */ + + size_t bias_batch_size; /*!< Bias broadcast dim for batch. */ + size_t bias_num_heads; /*!< Bias broadcast dim for heads. */ + size_t bias_seqlen_q; /*!< Bias broadcast dim for Q sequence length. */ + size_t bias_seqlen_kv; /*!< Bias broadcast dim for K/V sequence length. */ + + bool is_training; /*!< Whether the model is in training mode. */ + bool return_max_logit; /*!< Whether to produce Max along with Stats (fwd-only). */ + bool deterministic; /*!< Whether determinism is required (bwd-only). */ +} NVTEFusedAttnConfig; + +/*! \brief Default-initialize an ``NVTEFusedAttnConfig``. + * + * Sets ``struct_size`` and the categorical fields (layouts, formats, masks, + * window sizes, scaling mode) to safe NOT_SET / no-op defaults. Numeric and + * tensor-derived fields, paged-KV shape, bias broadcast shape, and direction + * flags all default to zero/false; callers must set the fields relevant to + * their query. + */ +#define NVTE_FUSED_ATTN_CONFIG_INIT \ + { \ + .struct_size = sizeof(NVTEFusedAttnConfig), \ + .qkv_layout = NVTE_QKV_Layout_NOT_SET, .o_format = NVTE_QKV_Format_NOT_SET, \ + .do_format = NVTE_QKV_Format_NOT_SET, .dqkv_layout = NVTE_QKV_Layout_NOT_SET, \ + .qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET, \ + .do_scale_inv_format = NVTE_QKV_Format_NOT_SET, .bias_type = NVTE_NO_BIAS, \ + .attn_mask_type = NVTE_NO_MASK, .softmax_type = NVTE_VANILLA_SOFTMAX, \ + .scaling_mode = NVTE_DELAYED_TENSOR_SCALING, .window_size_left = -1, .window_size_right = -1, \ + } + +/*! \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. Must be initialized + * with ``NVTE_FUSED_ATTN_CONFIG_INIT`` and have + * ``cfg->struct_size`` set to ``sizeof(NVTEFusedAttnConfig)``. + * \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(const NVTEFusedAttnConfig *cfg, + const char **message); + /*! \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] is_training Whether the model is in training mode. - * \param[in] batch_size Batch size. - * \param[in] q_dtype The data type of Tensor Q. - * \param[in] kv_dtype The data type of Tensors K, V. - * \param[in] o_dtype The data type of Tensor O. - * \param[in] scaling_mode Scaling mode of attention. - * \param[in] qkv_layout The layout of Tensors Q, K, V. - * \param[in] o_format The format of Tensor O. - * \param[in] do_format The format of Tensor dO. - * \param[in] dqkv_layout The layout of Tensors dQ, dK, dV. - * \param[in] qkv_scale_inv_format Format of the scale-inverse tensors for QKV in FP8 - * configurations; pass NVTE_QKV_Format_NOT_SET to let the - * backend infer it from ``qkv_layout`` otherwise. - * \param[in] do_scale_inv_format Format of the scale-inverse tensor for dO in FP8 backward - * configurations; pass NVTE_QKV_Format_NOT_SET to let the - * backend infer it from ``do_format`` otherwise. - * \param[in] bias_type The attention bias type. - * \param[in] attn_mask_type The attention mask type. - * \param[in] softmax_type The attention softmax type. - * \param[in] attn_scale Scaling factor for Q * K^T. - * \param[in] dropout The dropout probability. - * \param[in] num_attn_heads The number of heads in Q. - * \param[in] num_gqa_groups The number of heads in K, V. - * \param[in] max_seqlen_q The sequence length of Q. - * \param[in] max_seqlen_kv The sequence length of K, V. - * \param[in] head_dim_qk The head dimension of Q, K. - * \param[in] head_dim_v The head dimension of V. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the - * bottom right corner of the softmax matrix. - * \param[in] return_max_logit Whether to produce Max along with Stats. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] deterministic Whether determinism is required or not. - * \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`` - * on the same thread; callers that need to retain the - * message across further calls must copy it. Pass NULL to - * skip diagnostics. + * \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. + * \param[in] kv_dtype The data type of Tensors K, V. + * \param[in] qkv_layout The layout of Tensors Q, K, V. + * \param[in] bias_type The attention bias type. + * \param[in] attn_mask_type The attention mask type. + * \param[in] softmax_type The attention softmax type. + * \param[in] dropout The dropout probability. + * \param[in] num_attn_heads The number of heads in Q. + * \param[in] num_gqa_groups The number of heads in K, V. + * \param[in] max_seqlen_q The sequence length of Q. + * \param[in] max_seqlen_kv The sequence length of K, V. + * \param[in] head_dim_qk The head dimension of Q, K. + * \param[in] head_dim_v The head dimension of V. + * \param[in] window_size_left Sliding window size (the left half). + * \param[in] window_size_right Sliding window size (the right half). + * \param[in] return_max_logit Whether to produce Max along with Stats. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. + * \param[in] deterministic Whether determinism is required or not. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, size_t batch_size, NVTEDType q_dtype, NVTEDType kv_dtype, NVTEDType o_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 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, - const char **message); + 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); /*! \brief Compute dot product attention with separate Q, K and V. * diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index c00feec748..e24a5c4b1b 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -129,6 +129,7 @@ class FusedAttnHelper: 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. @@ -162,6 +163,7 @@ def get_fused_attn_backend(self): 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, @@ -380,6 +382,7 @@ def abstract( 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: diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index cf455357c6..e181f7ed5a 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -152,9 +152,9 @@ std::tuple GetFusedAttnBackend( 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 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, + NVTE_Softmax_Type softmax_type, 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( diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 9c6f2483cd..b5e40aaf6a 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -14,12 +14,13 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 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 bottom_right_diagonal, bool deterministic) { + 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 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); } @@ -29,15 +30,40 @@ std::tuple GetFusedAttnBackend( 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."); + + NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + 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 = mask_type; + cfg.softmax_type = softmax_type; + cfg.scaling_mode = scaling_mode; + cfg.attn_scale = attn_scale; + cfg.dropout = dropout_probability; + cfg.max_seqlen_q = q_max_seqlen; + cfg.max_seqlen_kv = kv_max_seqlen; + cfg.window_size_left = window_size_left; + cfg.window_size_right = window_size_right; + cfg.bottom_right_diagonal = bottom_right_diagonal; + cfg.cuda_graph = false; + cfg.qkv_dtype = static_cast(q_dtype); + cfg.o_dtype = static_cast(o_dtype); + cfg.batch_size = batch_size; + cfg.num_attn_heads = q_attn_heads; + cfg.num_gqa_groups = kv_attn_heads; + cfg.head_dim_qk = qk_head_dim; + cfg.head_dim_v = v_head_dim; + cfg.is_training = is_training; + cfg.return_max_logit = false; + cfg.deterministic = deterministic; + const char *message = nullptr; - auto backend = nvte_get_fused_attn_backend( - is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), - static_cast(o_dtype), scaling_mode, qkv_layout, o_format, do_format, dqkv_layout, - qkv_scale_inv_format, do_scale_inv_format, bias_type, mask_type, softmax_type, - /*attn_scale=*/1.0f, 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, - bottom_right_diagonal, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, - &message); + auto backend = nvte_get_fused_attn_backend_v2(&cfg, &message); return {backend, message != nullptr ? std::string(message) : std::string()}; } @@ -277,17 +303,13 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - const NVTE_QKV_Format probe_o_format = nvte_get_q_format(qkv_layout); - auto backend = nvte_get_fused_attn_backend( - is_training, input_batch, static_cast(dtype), static_cast(dtype), - static_cast(dtype), NVTE_INVALID_SCALING, qkv_layout, probe_o_format, - /*do_format=*/probe_o_format, /*dqkv_layout=*/qkv_layout, - /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, - /*do_scale_inv_format=*/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, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, - /*message=*/nullptr); + auto [backend, _fwd_msg] = GetFusedAttnBackend( + is_training, input_batch, dtype, dtype, dtype, NVTE_INVALID_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) */ @@ -559,19 +581,13 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); - // JAX uses the same layout for output / dQKV as for QKV, so derive the formats from qkv_layout. - // Scale-inv formats stay NOT_SET because JAX's fused-attn path here is non-FP8. - const NVTE_QKV_Format probe_o_format = nvte_get_q_format(qkv_layout); - auto backend = nvte_get_fused_attn_backend( - is_training, input_batch, static_cast(dtype), static_cast(dtype), - static_cast(dtype), NVTE_INVALID_SCALING, qkv_layout, probe_o_format, - /*do_format=*/probe_o_format, /*dqkv_layout=*/qkv_layout, - /*qkv_scale_inv_format=*/NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, - /*do_scale_inv_format=*/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, /*return_max_logit=*/false, /*cuda_graph=*/false, deterministic, - /*message=*/nullptr); + auto [backend, _bwd_msg] = GetFusedAttnBackend( + is_training, input_batch, dtype, dtype, dtype, NVTE_INVALID_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/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 17e9a337a4..4112e7c922 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 @@ -426,6 +426,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"))) @@ -1441,6 +1442,7 @@ def forward( return_max_logit=self.return_max_logit, cuda_graph=is_graph_capturing(), num_splits=num_splits, + softmax_scale=self.softmax_scale, ) global _attention_backends if is_in_onnx_export_mode(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index cf2f297083..169add4d6e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -257,6 +257,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. """ qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor @@ -290,6 +293,7 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 + softmax_scale: float = 1.0 def __eq__(self, other): """ @@ -368,6 +372,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 # Run config logger = logging.getLogger("DotProductAttention") @@ -1262,6 +1267,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt AttnBiasType[fu_core_attention_bias_type], AttnMaskType[attn_mask_type], SoftmaxType[softmax_type], + softmax_scale, attention_dropout, num_heads, num_gqa_groups, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 74021b81b5..1e2b3d356e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -83,10 +83,10 @@ std::tuple get_fused_attn_backend( 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 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); + 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 4cda724a8b..2eff5c21b7 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -46,18 +46,44 @@ std::tuple get_fused_attn_backend( 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 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) { + 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) { + NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + 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 = p_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.cuda_graph = cuda_graph; + NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + cfg.qkv_dtype = static_cast(q_dtype); + cfg.o_dtype = static_cast(o_dtype); + cfg.batch_size = batch_size; + 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 = is_training; + cfg.return_max_logit = return_max_logit; + cfg.deterministic = deterministic; + const char *message = nullptr; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, batch_size, static_cast(q_dtype), static_cast(kv_dtype), - static_cast(o_dtype), scaling_mode, qkv_layout, o_format, do_format, dqkv_layout, - qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, softmax_type, - /*attn_scale=*/1.0f, 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, bottom_right_diagonal, - return_max_logit, cuda_graph, deterministic, &message); + 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()}; } From e86fc6712bd30b71e5da474329d64639ae063d96 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 May 2026 21:44:13 +0000 Subject: [PATCH 20/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/fused_attn.cpp | 17 +++--- .../common/fused_attn/fused_attn_fp8.cu | 4 +- .../include/transformer_engine/fused_attn.h | 54 ++++++++++--------- .../jax/csrc/extensions/attention.cpp | 13 +++-- .../pytorch/csrc/extensions/attention.cpp | 3 +- 5 files changed, 47 insertions(+), 44 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 6670fd59ed..828f98b50f 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -244,7 +244,7 @@ void set_message(const char **message, std::string reason) { // select a backend for fused attention NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig *cfg, - const char **message) { + const char **message) { using namespace transformer_engine; set_message(message, ""); NVTE_CHECK(cfg != nullptr, "NVTEFusedAttnConfig pointer must not be NULL."); @@ -282,8 +282,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig 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); + const bool is_f16_or_bf16 = + (cfg->qkv_dtype == NVTEDType::kNVTEFloat16 || cfg->qkv_dtype == NVTEDType::kNVTEBFloat16); if (is_fp8) { if (cfg->return_max_logit) { @@ -336,8 +336,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_message(message, - "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg->qkv_dtype) + " ."); + set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg->qkv_dtype) + " ."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -475,8 +474,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; cfg.qkv_layout = qkv_layout; cfg.o_format = o_format; - cfg.do_format = o_format; // fwd path: same format used for dO if/when probed for bwd - cfg.dqkv_layout = qkv_layout; // fwd path: same layout used for dQKV if/when probed for bwd + cfg.do_format = o_format; // fwd path: same format used for dO if/when probed for bwd + cfg.dqkv_layout = qkv_layout; // fwd path: same layout used for dQKV if/when probed for bwd cfg.qkv_scale_inv_format = qkv_scale_inv_format; cfg.do_scale_inv_format = qkv_scale_inv_format; // fwd path: mirror QKV cfg.bias_type = bias_type; @@ -493,8 +492,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.cuda_graph = cuda_graph; cfg.qkv_dtype = Q_type; cfg.o_dtype = O_type; - cfg.do_dtype = O_type; // fwd path: dO assumed to share dtype with O - cfg.dqkv_dtype = Q_type; // fwd path: dQKV assumed to share dtype with QKV + cfg.do_dtype = O_type; // fwd path: dO assumed to share dtype with O + cfg.dqkv_dtype = Q_type; // fwd path: dQKV assumed to share dtype with QKV cfg.batch_size = b; cfg.num_attn_heads = h_q; cfg.num_gqa_groups = h_kv; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 3fe5b7fb10..be689f2b0c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1325,7 +1325,7 @@ void fused_attn_fp8_bwd( } } -std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { +std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t handle) { 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; @@ -1376,7 +1376,7 @@ std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t h } } -std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { +std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t handle) { 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; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index df15148350..dba7dd68d6 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -207,9 +207,9 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * in range according to ``struct_size`` and uses safe defaults otherwise. */ typedef struct NVTEFusedAttnConfig { - size_t struct_size; /*!< MUST equal sizeof(NVTEFusedAttnConfig). */ - uint32_t reserved0; /*!< Padding for layout stability; set to 0. */ - uint32_t reserved1; /*!< Padding for layout stability; set to 0. */ + size_t struct_size; /*!< MUST equal sizeof(NVTEFusedAttnConfig). */ + uint32_t reserved0; /*!< Padding for layout stability; set to 0. */ + uint32_t reserved1; /*!< Padding for layout stability; set to 0. */ NVTE_QKV_Layout qkv_layout; /*!< QKV tensors' layout. */ NVTE_QKV_Format o_format; /*!< Output O tensor format. */ @@ -227,18 +227,18 @@ typedef struct NVTEFusedAttnConfig { size_t max_seqlen_kv; /*!< Max sequence length for K, V. */ int64_t window_size_left; /*!< Sliding window size (left half); -1 = unlimited. */ int64_t window_size_right; /*!< Sliding window size (right half); -1 = unlimited. */ - bool bottom_right_diagonal; /*!< Whether causal mask aligns to the bottom-right diagonal. */ - bool cuda_graph; /*!< Whether CUDA graph capture is enabled. */ - - NVTEDType qkv_dtype; /*!< Data type of Tensors Q, K, V. Q and K/V must share a dtype. */ - NVTEDType o_dtype; /*!< Data type of Tensor O. */ - NVTEDType do_dtype; /*!< Data type of Tensor dO (bwd). */ - NVTEDType dqkv_dtype; /*!< Data type of Tensors dQ, dK, dV (bwd). */ - size_t batch_size; /*!< Batch size. */ - size_t num_attn_heads; /*!< Number of heads in Q. */ - size_t num_gqa_groups; /*!< Number of heads in K, V. */ - size_t head_dim_qk; /*!< Head dimension of Q, K. */ - size_t head_dim_v; /*!< Head dimension of V. */ + bool bottom_right_diagonal; /*!< Whether causal mask aligns to the bottom-right diagonal. */ + bool cuda_graph; /*!< Whether CUDA graph capture is enabled. */ + + NVTEDType qkv_dtype; /*!< Data type of Tensors Q, K, V. Q and K/V must share a dtype. */ + NVTEDType o_dtype; /*!< Data type of Tensor O. */ + NVTEDType do_dtype; /*!< Data type of Tensor dO (bwd). */ + NVTEDType dqkv_dtype; /*!< Data type of Tensors dQ, dK, dV (bwd). */ + size_t batch_size; /*!< Batch size. */ + size_t num_attn_heads; /*!< Number of heads in Q. */ + size_t num_gqa_groups; /*!< Number of heads in K, V. */ + size_t head_dim_qk; /*!< Head dimension of Q, K. */ + size_t head_dim_v; /*!< Head dimension of V. */ size_t num_pages_k; /*!< Total number of K cache pages. */ size_t num_pages_v; /*!< Total number of V cache pages. */ @@ -265,15 +265,21 @@ typedef struct NVTEFusedAttnConfig { * flags all default to zero/false; callers must set the fields relevant to * their query. */ -#define NVTE_FUSED_ATTN_CONFIG_INIT \ - { \ - .struct_size = sizeof(NVTEFusedAttnConfig), \ - .qkv_layout = NVTE_QKV_Layout_NOT_SET, .o_format = NVTE_QKV_Format_NOT_SET, \ - .do_format = NVTE_QKV_Format_NOT_SET, .dqkv_layout = NVTE_QKV_Layout_NOT_SET, \ - .qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET, \ - .do_scale_inv_format = NVTE_QKV_Format_NOT_SET, .bias_type = NVTE_NO_BIAS, \ - .attn_mask_type = NVTE_NO_MASK, .softmax_type = NVTE_VANILLA_SOFTMAX, \ - .scaling_mode = NVTE_DELAYED_TENSOR_SCALING, .window_size_left = -1, .window_size_right = -1, \ +#define NVTE_FUSED_ATTN_CONFIG_INIT \ + { \ + .struct_size = sizeof(NVTEFusedAttnConfig), \ + .qkv_layout = NVTE_QKV_Layout_NOT_SET, \ + .o_format = NVTE_QKV_Format_NOT_SET, \ + .do_format = NVTE_QKV_Format_NOT_SET, \ + .dqkv_layout = NVTE_QKV_Layout_NOT_SET, \ + .qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET, \ + .do_scale_inv_format = NVTE_QKV_Format_NOT_SET, \ + .bias_type = NVTE_NO_BIAS, \ + .attn_mask_type = NVTE_NO_MASK, \ + .softmax_type = NVTE_VANILLA_SOFTMAX, \ + .scaling_mode = NVTE_DELAYED_TENSOR_SCALING, \ + .window_size_left = -1, \ + .window_size_right = -1, \ } /*! \brief Get fused attention backend based on input parameters. diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index b5e40aaf6a..cd09628be0 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -14,13 +14,12 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 mask_type, NVTE_Softmax_Type softmax_type, - 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) { + 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 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); } diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 2eff5c21b7..84f72ff879 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -82,8 +82,7 @@ std::tuple get_fused_attn_backend( cfg.deterministic = deterministic; const char *message = nullptr; - NVTE_Fused_Attn_Backend fused_attention_backend = - nvte_get_fused_attn_backend_v2(&cfg, &message); + 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()}; } From e2561d0d4c406b7bef4dacf8b9c7abf9db1e6c87 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 12 May 2026 16:20:35 -0700 Subject: [PATCH 21/63] fix FP8 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 4c8435f246..681dbea2c8 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1775,12 +1775,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: @@ -2567,6 +2578,7 @@ 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 From 724a12f009d96a0dc63775e5cff8cae6023d0c33 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 12 May 2026 16:20:52 -0700 Subject: [PATCH 22/63] add do_dtype and dqkv_dtype to API Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 15 +++++-------- .../common/fused_attn/fused_attn_fp8.cu | 6 ++++-- .../jax/cpp_extensions/attention.py | 2 ++ transformer_engine/jax/csrc/extensions.h | 15 ++++++------- .../jax/csrc/extensions/attention.cpp | 21 +++++++++++-------- .../attention/dot_product_attention/utils.py | 10 +++++++++ transformer_engine/pytorch/csrc/extensions.h | 5 +++-- .../pytorch/csrc/extensions/attention.cpp | 7 +++++-- 8 files changed, 49 insertions(+), 32 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 828f98b50f..0fbd9a21ae 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -342,17 +342,17 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig // 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, etc.) -// that this wrapper cannot express. +// 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; NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; cfg.qkv_layout = qkv_layout; - cfg.dqkv_layout = qkv_layout; // legacy: gradient layout matches input layout cfg.bias_type = bias_type; cfg.attn_mask_type = attn_mask_type; cfg.softmax_type = softmax_type; @@ -371,7 +371,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( cfg.num_gqa_groups = num_gqa_groups; cfg.head_dim_qk = head_dim_qk; cfg.head_dim_v = head_dim_v; - cfg.is_training = is_training; + 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 nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); @@ -474,10 +474,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; cfg.qkv_layout = qkv_layout; cfg.o_format = o_format; - cfg.do_format = o_format; // fwd path: same format used for dO if/when probed for bwd - cfg.dqkv_layout = qkv_layout; // fwd path: same layout used for dQKV if/when probed for bwd cfg.qkv_scale_inv_format = qkv_scale_inv_format; - cfg.do_scale_inv_format = qkv_scale_inv_format; // fwd path: mirror QKV cfg.bias_type = bias_type; cfg.attn_mask_type = attn_mask_type; cfg.softmax_type = softmax_type; @@ -492,8 +489,6 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.cuda_graph = cuda_graph; cfg.qkv_dtype = Q_type; cfg.o_dtype = O_type; - cfg.do_dtype = O_type; // fwd path: dO assumed to share dtype with O - cfg.dqkv_dtype = Q_type; // fwd path: dQKV assumed to share dtype with QKV cfg.batch_size = b; cfg.num_attn_heads = h_q; cfg.num_gqa_groups = h_kv; @@ -509,7 +504,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.bias_num_heads = bias_h; cfg.bias_seqlen_q = bias_sq; cfg.bias_seqlen_kv = bias_skv; - cfg.is_training = is_training; + cfg.is_training = false; cfg.return_max_logit = return_max_logit; cfg.deterministic = false; NVTE_Fused_Attn_Backend fused_attention_backend = diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index be689f2b0c..180bee2ab0 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1401,12 +1401,14 @@ std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t h const bool deterministic = cfg->deterministic; const DType qkv_dtype = static_cast(cfg->qkv_dtype); const DType o_dtype = static_cast(cfg->o_dtype); + const DType do_dtype = static_cast(cfg->do_dtype); + const DType dqkv_dtype = static_cast(cfg->dqkv_dtype); const NVTEScalingMode scaling_mode = cfg->scaling_mode; const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); - const cudnn_frontend::DataType_t do_t = o_t; - const cudnn_frontend::DataType_t dqkv_t = qkv_t; + const cudnn_frontend::DataType_t do_t = get_cudnn_fe_dtype(do_dtype); + const cudnn_frontend::DataType_t dqkv_t = get_cudnn_fe_dtype(dqkv_dtype); size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_bwd_impl( diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index e24a5c4b1b..a895d8eac3 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -153,6 +153,8 @@ def get_fused_attn_backend(self): q_type, jax_dtype_to_te_dtype(self.kv_dtype), q_type, + q_type, + q_type, NVTEScalingMode.NVTE_INVALID_SCALING, self.qkv_layout.value, NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index e181f7ed5a..b2adb3b042 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -149,13 +149,14 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 mask_type, - NVTE_Softmax_Type softmax_type, 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); + DType do_dtype, 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 mask_type, NVTE_Softmax_Type softmax_type, + 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 cd09628be0..573186b78d 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -13,13 +13,14 @@ namespace jax { std::tuple GetFusedAttnBackend( bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_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 mask_type, - NVTE_Softmax_Type softmax_type, 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) { + DType do_dtype, 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 mask_type, NVTE_Softmax_Type softmax_type, + 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); } @@ -52,6 +53,8 @@ std::tuple GetFusedAttnBackend( cfg.cuda_graph = false; cfg.qkv_dtype = static_cast(q_dtype); cfg.o_dtype = static_cast(o_dtype); + cfg.do_dtype = static_cast(do_dtype); + cfg.dqkv_dtype = static_cast(dqkv_dtype); cfg.batch_size = batch_size; cfg.num_attn_heads = q_attn_heads; cfg.num_gqa_groups = kv_attn_heads; @@ -303,7 +306,7 @@ static void FusedAttnForwardImpl( auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); auto [backend, _fwd_msg] = GetFusedAttnBackend( - is_training, input_batch, dtype, dtype, dtype, NVTE_INVALID_SCALING, qkv_layout, + is_training, input_batch, dtype, dtype, dtype, dtype, dtype, NVTE_INVALID_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, @@ -581,7 +584,7 @@ static void FusedAttnBackwardImpl( NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); auto [backend, _bwd_msg] = GetFusedAttnBackend( - is_training, input_batch, dtype, dtype, dtype, NVTE_INVALID_SCALING, qkv_layout, + is_training, input_batch, dtype, dtype, dtype, dtype, dtype, NVTE_INVALID_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, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 169add4d6e..a345e2c352 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1229,6 +1229,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt 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 @@ -1240,14 +1242,20 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt 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 @@ -1257,6 +1265,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt q_type, kv_type, o_type, + do_type, + dqkv_type, scaling_mode, QKVLayout[qkv_layout], QKVFormat[o_format], diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1e2b3d356e..019bf5afea 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -79,8 +79,9 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T // 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, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + 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, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 84f72ff879..3f2a1d4399 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -42,8 +42,9 @@ namespace transformer_engine::pytorch { // get the 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, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + 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, @@ -72,6 +73,8 @@ std::tuple get_fused_attn_backend( NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); cfg.qkv_dtype = static_cast(q_dtype); cfg.o_dtype = static_cast(o_dtype); + cfg.do_dtype = static_cast(do_dtype); + cfg.dqkv_dtype = static_cast(dqkv_dtype); cfg.batch_size = batch_size; cfg.num_attn_heads = num_attn_heads; cfg.num_gqa_groups = num_gqa_groups; From 3ae36df37a92a370a57943da30a781ff2696e6c4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 12 May 2026 23:24:23 +0000 Subject: [PATCH 23/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/extensions.h | 17 ++++++++--------- .../pytorch/csrc/extensions/attention.cpp | 17 ++++++++--------- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 019bf5afea..9931d4b3a8 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -79,15 +79,14 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T // 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); + 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 3f2a1d4399..afcdcae015 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -42,15 +42,14 @@ namespace transformer_engine::pytorch { // get the 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) { + 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) { NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; cfg.qkv_layout = qkv_layout; cfg.o_format = o_format; From 1c090726a50c0bb584edadd5e8926004df92f513 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Sat, 11 Jul 2026 08:37:25 -0700 Subject: [PATCH 24/63] replace with opaque handle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/common/CMakeLists.txt | 1 + .../common/fused_attn/config_and_params.cpp | 423 +++++++++++++++++ .../common/fused_attn/config_and_params.h | 161 +++++++ .../common/fused_attn/fused_attn.cpp | 155 ++++--- .../fused_attn_f16_arbitrary_seqlen.cu | 424 +++++++----------- .../fused_attn_f16_arbitrary_seqlen.h | 34 +- .../common/fused_attn/fused_attn_fp8.cu | 352 +++++---------- .../common/fused_attn/fused_attn_fp8.h | 31 +- transformer_engine/common/fused_attn/utils.cu | 40 ++ transformer_engine/common/fused_attn/utils.h | 60 --- .../include/transformer_engine/fused_attn.h | 415 +++++++++++++---- .../jax/cpp_extensions/attention.py | 4 +- transformer_engine/jax/csrc/extensions.h | 2 +- .../jax/csrc/extensions/attention.cpp | 72 +-- .../jax/csrc/extensions/pybind.cpp | 8 - .../pytorch/csrc/extensions/attention.cpp | 65 +-- 16 files changed, 1400 insertions(+), 847 deletions(-) create mode 100644 transformer_engine/common/fused_attn/config_and_params.cpp create mode 100644 transformer_engine/common/fused_attn/config_and_params.h 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 0fbd9a21ae..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" @@ -243,26 +244,24 @@ void set_message(const char **message, std::string reason) { } // namespace // select a backend for fused attention -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig *cfg, - const char **message) { +namespace { + +NVTE_Fused_Attn_Backend select_fused_attn_backend(const transformer_engine::FusedAttnConfig &cfg, + const char **message) { using namespace transformer_engine; set_message(message, ""); - NVTE_CHECK(cfg != nullptr, "NVTEFusedAttnConfig pointer must not be NULL."); - NVTE_CHECK(cfg->struct_size == sizeof(NVTEFusedAttnConfig), - "NVTEFusedAttnConfig::struct_size must equal sizeof(NVTEFusedAttnConfig); " - "did you forget NVTE_FUSED_ATTN_CONFIG_INIT?"); 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 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(); // 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, 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); + 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 " @@ -272,21 +271,21 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig // 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) { + 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_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); + (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); if (is_fp8) { - if (cfg->return_max_logit) { + 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; } @@ -301,7 +300,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig set_message(message, std::move(fwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg->is_training) { + 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)); @@ -312,12 +311,12 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig } if (is_f16_or_bf16) { - if (cudnn_runtime_version <= 91500 && cfg->is_training && + if (cudnn_runtime_version <= 91500 && cfg.is_training && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - (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) { + (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; } @@ -326,7 +325,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig set_message(message, std::move(fwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg->is_training) { + 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)); @@ -336,10 +335,18 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg->qkv_dtype) + " ."); + 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, @@ -351,7 +358,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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; - NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + 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; @@ -374,7 +381,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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 nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); + return select_fused_attn_backend(cfg, /*message=*/nullptr); } // NVTE fused attention FWD with separate Q, K and V @@ -464,14 +471,19 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTEScalingMode scaling_mode = input_Q->scaling_mode; size_t bias_b = 0, bias_h = 0, bias_sq = 0, bias_skv = 0; - if (input_Bias->data.dptr != nullptr && input_Bias->data.shape.size() >= 4) { + 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]; } - NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + 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; @@ -486,7 +498,6 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.window_size_left = window_size_left; cfg.window_size_right = window_size_right; cfg.bottom_right_diagonal = bottom_right_diagonal; - cfg.cuda_graph = cuda_graph; cfg.qkv_dtype = Q_type; cfg.o_dtype = O_type; cfg.batch_size = b; @@ -504,28 +515,23 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.bias_num_heads = bias_h; cfg.bias_seqlen_q = bias_sq; cfg.bias_seqlen_kv = bias_skv; - cfg.is_training = false; - cfg.return_max_logit = return_max_logit; - cfg.deterministic = false; + cfg.num_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; NVTE_Fused_Attn_Backend fused_attention_backend = - nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); + 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"); } @@ -591,7 +597,20 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTEDType dQKV_type = static_cast(output_dQ->data.dtype); const NVTEScalingMode scaling_mode = input_Q->scaling_mode; - NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; + 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]; + } + + 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; @@ -609,7 +628,6 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.window_size_left = window_size_left; cfg.window_size_right = window_size_right; cfg.bottom_right_diagonal = bottom_right_diagonal; - cfg.cuda_graph = cuda_graph; cfg.qkv_dtype = Q_type; cfg.o_dtype = O_type; cfg.do_dtype = dO_type; @@ -619,11 +637,14 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cfg.num_gqa_groups = h_kv; cfg.head_dim_qk = d_qk; cfg.head_dim_v = d_v; - cfg.is_training = true; - cfg.return_max_logit = false; - cfg.deterministic = deterministic; + 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 = - nvte_get_fused_attn_backend_v2(&cfg, /*message=*/nullptr); + select_fused_attn_backend(cfg, /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; @@ -636,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++]); @@ -656,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 9cdee256ed..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) { @@ -1334,66 +1293,20 @@ void fused_attn_arbitrary_seqlen_bwd( } } -std::string is_supported_f16_fwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { - const size_t num_gqa_groups = cfg->num_gqa_groups; - const size_t head_dim_qk = cfg->head_dim_qk; - const size_t head_dim_v = cfg->head_dim_v; - const bool is_training = cfg->is_training; - const bool return_max_logit = cfg->return_max_logit; - const float attn_scale = cfg->attn_scale; - const float p_dropout = 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 DType qkv_dtype = static_cast(cfg->qkv_dtype); - const auto b = static_cast(cfg->batch_size); - const auto h = static_cast(cfg->num_attn_heads); - const auto sq = static_cast(cfg->max_seqlen_q); - const auto skv = static_cast(cfg->max_seqlen_kv); - - const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - const NVTE_QKV_Format kv_format = nvte_get_kv_format(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 NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - const bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); - const bool has_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - - const int64_t max_b = (is_ragged_q || is_ragged_kv) ? b : 0; - const int64_t max_t_q = is_ragged_q ? b * sq : 0; - const int64_t max_t_kv = is_ragged_kv ? b * skv : 0; - const int64_t num_pages_k = is_paged_kv ? b : 0; - const int64_t num_pages_v = is_paged_kv ? b : 0; - const int64_t page_size_k = is_paged_kv ? skv : 0; - const int64_t page_size_v = is_paged_kv ? skv : 0; - const int64_t max_pages_per_seq_k = is_paged_kv ? 1 : 0; - const int64_t max_pages_per_seq_v = is_paged_kv ? 1 : 0; - const int64_t bias_b = has_bias ? b : 0; - const int64_t bias_h = has_bias ? h : 0; - const int64_t bias_sq = has_bias ? sq : 0; - const int64_t bias_skv = has_bias ? skv : 0; +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( - b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), - static_cast(head_dim_v), max_b, max_t_q, max_t_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, + 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, - get_cudnn_fe_dtype(qkv_dtype), /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; @@ -1404,51 +1317,15 @@ std::string is_supported_f16_fwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t h } } -std::string is_supported_f16_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t handle) { - const size_t num_gqa_groups = cfg->num_gqa_groups; - const size_t head_dim_qk = cfg->head_dim_qk; - const size_t head_dim_v = cfg->head_dim_v; - const float attn_scale = cfg->attn_scale; - const float p_dropout = 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 DType qkv_dtype = static_cast(cfg->qkv_dtype); - const auto b = static_cast(cfg->batch_size); - const auto h = static_cast(cfg->num_attn_heads); - const auto sq = static_cast(cfg->max_seqlen_q); - const auto skv = static_cast(cfg->max_seqlen_kv); - - const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - const NVTE_QKV_Format kv_format = nvte_get_kv_format(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 bool has_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - - const int64_t max_b = (is_ragged_q || is_ragged_kv) ? b : 0; - const int64_t max_t_q = is_ragged_q ? b * sq : 0; - const int64_t max_t_kv = is_ragged_kv ? b * skv : 0; - const int64_t bias_b = has_bias ? b : 0; - const int64_t bias_h = has_bias ? h : 0; - const int64_t bias_sq = has_bias ? sq : 0; - const int64_t bias_skv = has_bias ? skv : 0; +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( - b, h, static_cast(num_gqa_groups), sq, skv, static_cast(head_dim_qk), - static_cast(head_dim_v), max_b, max_t_q, max_t_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=*/nullptr, /*devPtrKTranspose=*/nullptr, + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrKTranspose=*/nullptr, /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, /*devPtrdO=*/nullptr, @@ -1456,7 +1333,6 @@ std::string is_supported_f16_bwd(const NVTEFusedAttnConfig *cfg, cudnnHandle_t h /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, - get_cudnn_fe_dtype(qkv_dtype), /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; 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 5d27e82278..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 @@ -16,33 +16,21 @@ #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, @@ -52,12 +40,12 @@ void fused_attn_arbitrary_seqlen_bwd( // 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 NVTEFusedAttnConfig *cfg, cudnnHandle_t handle); +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 NVTEFusedAttnConfig *cfg, cudnnHandle_t handle); +std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 180bee2ab0..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"); } @@ -1325,47 +1280,18 @@ void fused_attn_fp8_bwd( } } -std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t handle) { - 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 bool is_training = cfg->is_training; - const float attn_scale = cfg->attn_scale; - const float p_dropout = cfg->dropout; - const NVTE_QKV_Layout qkv_layout = cfg->qkv_layout; - const NVTE_QKV_Format o_format = cfg->o_format; - const NVTE_QKV_Format qkv_scale_inv_format = cfg->qkv_scale_inv_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 DType qkv_dtype = static_cast(cfg->qkv_dtype); - const DType o_dtype = static_cast(cfg->o_dtype); - const NVTEScalingMode scaling_mode = cfg->scaling_mode; - +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( - static_cast(batch), static_cast(num_attn_heads), - static_cast(num_gqa_groups), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(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, + 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, get_cudnn_fe_dtype(qkv_dtype), get_cudnn_fe_dtype(o_dtype), - scaling_mode, qkv_scale_inv_format, + /*devPtrDropoutOffset=*/nullptr, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; @@ -1376,48 +1302,11 @@ std::string is_supported_fp8_fwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t h } } -std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t handle) { - 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 float attn_scale = cfg->attn_scale; - const float p_dropout = 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_QKV_Format qkv_scale_inv_format = cfg->qkv_scale_inv_format; - const NVTE_QKV_Format do_scale_inv_format = cfg->do_scale_inv_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 bool deterministic = cfg->deterministic; - const DType qkv_dtype = static_cast(cfg->qkv_dtype); - const DType o_dtype = static_cast(cfg->o_dtype); - const DType do_dtype = static_cast(cfg->do_dtype); - const DType dqkv_dtype = static_cast(cfg->dqkv_dtype); - const NVTEScalingMode scaling_mode = cfg->scaling_mode; - - const cudnn_frontend::DataType_t qkv_t = get_cudnn_fe_dtype(qkv_dtype); - const cudnn_frontend::DataType_t o_t = get_cudnn_fe_dtype(o_dtype); - const cudnn_frontend::DataType_t do_t = get_cudnn_fe_dtype(do_dtype); - const cudnn_frontend::DataType_t dqkv_t = get_cudnn_fe_dtype(dqkv_dtype); +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( - static_cast(batch), static_cast(num_attn_heads), - static_cast(num_gqa_groups), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), static_cast(head_dim_qk), - static_cast(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, + cfg, /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, @@ -1431,8 +1320,7 @@ std::string is_supported_fp8_bwd(const NVTEFusedAttnConfig* cfg, cudnnHandle_t h /*devPtrdO_t=*/nullptr, /*devPtrDescaleQ_t=*/nullptr, /*devPtrDescaleK_t=*/nullptr, /*devPtrDescaledO_t=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, - /*devPtrDropoutOffset=*/nullptr, qkv_t, o_t, do_t, dqkv_t, scaling_mode, - qkv_scale_inv_format, do_scale_inv_format, + /*devPtrDropoutOffset=*/nullptr, /*workspace=*/nullptr, &workspace_size, /*stream=*/static_cast(0), handle); return ""; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index fc60987cf3..1ede20c7f1 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -10,45 +10,34 @@ #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 NVTEFusedAttnConfig *cfg, cudnnHandle_t handle); +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 NVTEFusedAttnConfig *cfg, cudnnHandle_t handle); +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 338277e904..f26e03c5ad 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -196,91 +196,88 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); -/*! \struct NVTEFusedAttnConfig - * \brief Attention configuration. - * - * Versioning rules: - * - ``struct_size`` MUST be set to ``sizeof(NVTEFusedAttnConfig)`` by the - * caller (use ``NVTE_FUSED_ATTN_CONFIG_INIT``). - * - New fields may only be appended at the end; existing fields are never - * reordered, removed, or resized. The library reads only fields that are - * in range according to ``struct_size`` and uses safe defaults otherwise. - */ -typedef struct NVTEFusedAttnConfig { - size_t struct_size; /*!< MUST equal sizeof(NVTEFusedAttnConfig). */ - uint32_t reserved0; /*!< Padding for layout stability; set to 0. */ - uint32_t reserved1; /*!< Padding for layout stability; set to 0. */ - - NVTE_QKV_Layout qkv_layout; /*!< QKV tensors' layout. */ - NVTE_QKV_Format o_format; /*!< Output O tensor format. */ - NVTE_QKV_Format do_format; /*!< Output-grad dO tensor format (bwd). */ - NVTE_QKV_Layout dqkv_layout; /*!< Gradient dQKV tensor layout (bwd). */ - NVTE_QKV_Format qkv_scale_inv_format; /*!< QKV scale_inv tensor format (FP8). */ - NVTE_QKV_Format do_scale_inv_format; /*!< dO scale_inv tensor format (FP8 bwd). */ - NVTE_Bias_Type bias_type; /*!< Attention bias type. */ - NVTE_Mask_Type attn_mask_type; /*!< Attention mask type. */ - NVTE_Softmax_Type softmax_type; /*!< Attention softmax type. */ - NVTEScalingMode scaling_mode; /*!< Scaling mode (e.g. delayed, MXFP8). */ - float attn_scale; /*!< Pre-softmax attention scale factor. */ - float dropout; /*!< Dropout probability. */ - size_t max_seqlen_q; /*!< Max sequence length for Q. */ - size_t max_seqlen_kv; /*!< Max sequence length for K, V. */ - int64_t window_size_left; /*!< Sliding window size (left half); -1 = unlimited. */ - int64_t window_size_right; /*!< Sliding window size (right half); -1 = unlimited. */ - bool bottom_right_diagonal; /*!< Whether causal mask aligns to the bottom-right diagonal. */ - bool cuda_graph; /*!< Whether CUDA graph capture is enabled. */ - - NVTEDType qkv_dtype; /*!< Data type of Tensors Q, K, V. Q and K/V must share a dtype. */ - NVTEDType o_dtype; /*!< Data type of Tensor O. */ - NVTEDType do_dtype; /*!< Data type of Tensor dO (bwd). */ - NVTEDType dqkv_dtype; /*!< Data type of Tensors dQ, dK, dV (bwd). */ - size_t batch_size; /*!< Batch size. */ - size_t num_attn_heads; /*!< Number of heads in Q. */ - size_t num_gqa_groups; /*!< Number of heads in K, V. */ - size_t head_dim_qk; /*!< Head dimension of Q, K. */ - size_t head_dim_v; /*!< Head dimension of V. */ - - size_t num_pages_k; /*!< Total number of K cache pages. */ - size_t num_pages_v; /*!< Total number of V cache pages. */ - size_t page_size_k; /*!< Tokens per K cache page. */ - size_t page_size_v; /*!< Tokens per V cache page. */ - size_t max_pages_per_seq_k; /*!< Max K pages per sequence in the batch. */ - size_t max_pages_per_seq_v; /*!< Max V pages per sequence in the batch. */ - - size_t bias_batch_size; /*!< Bias broadcast dim for batch. */ - size_t bias_num_heads; /*!< Bias broadcast dim for heads. */ - size_t bias_seqlen_q; /*!< Bias broadcast dim for Q sequence length. */ - size_t bias_seqlen_kv; /*!< Bias broadcast dim for K/V sequence length. */ - - bool is_training; /*!< Whether the model is in training mode. */ - bool return_max_logit; /*!< Whether to produce Max along with Stats (fwd-only). */ - bool deterministic; /*!< Whether determinism is required (bwd-only). */ -} NVTEFusedAttnConfig; - -/*! \brief Default-initialize an ``NVTEFusedAttnConfig``. - * - * Sets ``struct_size`` and the categorical fields (layouts, formats, masks, - * window sizes, scaling mode) to safe NOT_SET / no-op defaults. Numeric and - * tensor-derived fields, paged-KV shape, bias broadcast shape, and direction - * flags all default to zero/false; callers must set the fields relevant to - * their query. - */ -#define NVTE_FUSED_ATTN_CONFIG_INIT \ - { \ - .struct_size = sizeof(NVTEFusedAttnConfig), \ - .qkv_layout = NVTE_QKV_Layout_NOT_SET, \ - .o_format = NVTE_QKV_Format_NOT_SET, \ - .do_format = NVTE_QKV_Format_NOT_SET, \ - .dqkv_layout = NVTE_QKV_Layout_NOT_SET, \ - .qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET, \ - .do_scale_inv_format = NVTE_QKV_Format_NOT_SET, \ - .bias_type = NVTE_NO_BIAS, \ - .attn_mask_type = NVTE_NO_MASK, \ - .softmax_type = NVTE_VANILLA_SOFTMAX, \ - .scaling_mode = NVTE_DELAYED_TENSOR_SCALING, \ - .window_size_left = -1, \ - .window_size_right = -1, \ - } +/*! \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. * @@ -290,9 +287,9 @@ typedef struct NVTEFusedAttnConfig { * ``nvte_fused_attn_bwd`` to maintain a consistent signature between graph * building and runtime calls. * - * \param[in] cfg Attention configuration. Must be initialized - * with ``NVTE_FUSED_ATTN_CONFIG_INIT`` and have - * ``cfg->struct_size`` set to ``sizeof(NVTEFusedAttnConfig)``. + * \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 @@ -303,7 +300,7 @@ typedef struct NVTEFusedAttnConfig { * * \return Backend able to execute this configuration, or ``NVTE_No_Backend`` if none. */ -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(const NVTEFusedAttnConfig *cfg, +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, const char **message); /*! \brief Get fused attention backend based on input parameters. @@ -822,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/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index a895d8eac3..eaa9c8769a 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -17,10 +17,10 @@ import transformer_engine_jax from transformer_engine_jax import ( + JAXX_Scaling_Mode, NVTE_Fused_Attn_Backend, NVTE_QKV_Format, NVTE_QKV_Layout, - NVTEScalingMode, ) from transformer_engine.jax.attention import ( AttnBiasType, @@ -155,7 +155,7 @@ def get_fused_attn_backend(self): q_type, q_type, q_type, - NVTEScalingMode.NVTE_INVALID_SCALING, + 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, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 8769c3b8bd..5a6790793c 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -154,7 +154,7 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); 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, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, + 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, diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 25e59620da..c88f63a5e9 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -26,7 +26,7 @@ namespace jax { 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, NVTEScalingMode scaling_mode, NVTE_QKV_Layout qkv_layout, + 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, @@ -45,40 +45,40 @@ std::tuple GetFusedAttnBackend( } NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); - NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; - 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 = mask_type; - cfg.softmax_type = softmax_type; - cfg.scaling_mode = scaling_mode; - cfg.attn_scale = attn_scale; - cfg.dropout = dropout_probability; - cfg.max_seqlen_q = q_max_seqlen; - cfg.max_seqlen_kv = kv_max_seqlen; - cfg.window_size_left = window_size_left; - cfg.window_size_right = window_size_right; - cfg.bottom_right_diagonal = bottom_right_diagonal; - cfg.cuda_graph = false; - cfg.qkv_dtype = static_cast(q_dtype); - cfg.o_dtype = static_cast(o_dtype); - cfg.do_dtype = static_cast(do_dtype); - cfg.dqkv_dtype = static_cast(dqkv_dtype); - cfg.batch_size = batch_size; - cfg.num_attn_heads = q_attn_heads; - cfg.num_gqa_groups = kv_attn_heads; - cfg.head_dim_qk = qk_head_dim; - cfg.head_dim_v = v_head_dim; - cfg.is_training = is_training; - cfg.return_max_logit = false; - cfg.deterministic = deterministic; + 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); + auto backend = nvte_get_fused_attn_backend_v2(cfg, &message); return {backend, message != nullptr ? std::string(message) : std::string()}; } @@ -319,7 +319,8 @@ static void FusedAttnForwardImpl( auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); auto [backend, _fwd_msg] = GetFusedAttnBackend( - is_training, input_batch, dtype, dtype, dtype, dtype, dtype, NVTE_INVALID_SCALING, qkv_layout, + 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, @@ -597,7 +598,8 @@ static void FusedAttnBackwardImpl( NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); auto [backend, _bwd_msg] = GetFusedAttnBackend( - is_training, input_batch, dtype, dtype, dtype, dtype, dtype, NVTE_INVALID_SCALING, qkv_layout, + 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, diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index a4500994e8..bfb7b2a826 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -248,14 +248,6 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVFP4_2D_SCALING", JAXX_Scaling_Mode::NVFP4_2D_SCALING) .export_values(); - 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, "JAXX_Quantize_Layout", pybind11::module_local()) .value("ROWWISE", JAXX_Quantize_Layout::ROWWISE) .value("COLWISE", JAXX_Quantize_Layout::COLWISE) diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 9265dfc708..706eb630e0 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -50,41 +50,42 @@ std::tuple get_fused_attn_backend( 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) { - NVTEFusedAttnConfig cfg = NVTE_FUSED_ATTN_CONFIG_INIT; - 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 = p_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.cuda_graph = cuda_graph; NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); - cfg.qkv_dtype = static_cast(q_dtype); - cfg.o_dtype = static_cast(o_dtype); - cfg.do_dtype = static_cast(do_dtype); - cfg.dqkv_dtype = static_cast(dqkv_dtype); - cfg.batch_size = batch_size; - 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 = is_training; - cfg.return_max_logit = return_max_logit; - cfg.deterministic = deterministic; + + 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); + 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()}; } From ac19f9d070aa295251770fd2e3c0d3646d9d5121 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:07:45 -0700 Subject: [PATCH 25/63] repeat with fwd/bwd params Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.cpp | 812 +++++++++++++++--- .../common/fused_attn/config_and_params.h | 317 +++++-- .../common/fused_attn/fused_attn.cpp | 338 ++++---- .../fused_attn_f16_arbitrary_seqlen.cu | 49 +- .../common/fused_attn/fused_attn_fp8.cu | 14 +- transformer_engine/common/fused_attn/utils.cu | 37 - .../include/transformer_engine/fused_attn.h | 529 +++++++++++- .../dot_product_attention.py | 39 +- .../attention/dot_product_attention/utils.py | 215 +++-- transformer_engine/pytorch/csrc/extensions.h | 10 +- .../pytorch/csrc/extensions/attention.cpp | 88 +- .../pytorch/csrc/extensions/pybind.cpp | 2 +- 12 files changed, 1919 insertions(+), 531 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index c1e44e80af..79976fab14 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -6,8 +6,12 @@ #include "config_and_params.h" +#include + #include +#include "../util/cuda_runtime.h" + namespace { void bool_to_uint8(bool in, void *out) { @@ -94,6 +98,95 @@ void populate_fused_attn_config(FusedAttnConfig *cfg) { } } +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; + } + } + + // cuDNN graph supports dynamic shapes for batch_size + cache_cfg.batch_size = 1; + cache_cfg.bucketed_batch_size = 1; + cache_cfg.attention_scale = 1.0f; + + return cache_cfg; +} + +FusedAttnConfig make_fused_attn_config(const FusedAttnFwdParams ¶ms) { + FusedAttnConfig cfg = make_default_fused_attn_config(); + cfg.is_training = false; // fwd-only probe; caller restores before dispatch + cfg.deterministic = false; + cfg.cuda_graph = params.cuda_graph; + cfg.return_max_logit = params.return_max_logit; + cfg.qkv_layout = params.qkv_layout; + cfg.o_format = params.o_format; + cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; + cfg.bias_type = params.bias_type; + cfg.attn_mask_type = params.attn_mask_type; + cfg.softmax_type = params.softmax_type; + cfg.attn_scale = params.attn_scale; + cfg.dropout = params.dropout; + cfg.max_seqlen_q = params.max_seqlen_q; + cfg.max_seqlen_kv = params.max_seqlen_kv; + cfg.window_size_left = params.window_size_left; + cfg.window_size_right = params.window_size_right; + cfg.bottom_right_diagonal = params.bottom_right_diagonal; + return cfg; +} + +FusedAttnConfig make_fused_attn_config(const FusedAttnBwdParams ¶ms) { + FusedAttnConfig cfg = make_default_fused_attn_config(); + cfg.is_training = true; + cfg.deterministic = params.deterministic; + cfg.cuda_graph = params.cuda_graph; + cfg.return_max_logit = false; + cfg.qkv_layout = params.qkv_layout; + cfg.o_format = params.o_format; + cfg.do_format = params.do_format; + cfg.dqkv_layout = params.dqkv_layout; + cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; + cfg.do_scale_inv_format = params.do_scale_inv_format; + cfg.bias_type = params.bias_type; + cfg.attn_mask_type = params.attn_mask_type; + cfg.softmax_type = params.softmax_type; + cfg.attn_scale = params.attn_scale; + cfg.dropout = params.dropout; + cfg.max_seqlen_q = params.max_seqlen_q; + cfg.max_seqlen_kv = params.max_seqlen_kv; + cfg.window_size_left = params.window_size_left; + cfg.window_size_right = params.window_size_right; + cfg.bottom_right_diagonal = params.bottom_right_diagonal; + return cfg; +} + } // namespace transformer_engine NVTEFusedAttnConfig nvte_create_fused_attn_config() { @@ -138,29 +231,20 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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); + case kNVTEFusedAttnConfigAttnMaskType: + std::memcpy(buf, &cfg.attn_mask_type, 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); + 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 kNVTEFusedAttnConfigSoftmaxType: std::memcpy(buf, &cfg.softmax_type, attr_size); @@ -168,27 +252,9 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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; @@ -201,6 +267,27 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigDQKVDtype: std::memcpy(buf, &cfg.dqkv_dtype, attr_size); 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 kNVTEFusedAttnConfigAttnScale: + std::memcpy(buf, &cfg.attn_scale, attr_size); + break; case kNVTEFusedAttnConfigBatchSize: std::memcpy(buf, &cfg.batch_size, attr_size); break; @@ -216,6 +303,27 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigHeadDimV: std::memcpy(buf, &cfg.head_dim_v, 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 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; case kNVTEFusedAttnConfigNumPagesK: std::memcpy(buf, &cfg.num_pages_k, attr_size); break; @@ -246,21 +354,6 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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), ")"); } @@ -294,29 +387,20 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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); + case kNVTEFusedAttnConfigAttnMaskType: + std::memcpy(&cfg.attn_mask_type, 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); + 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 kNVTEFusedAttnConfigSoftmaxType: std::memcpy(&cfg.softmax_type, buf, attr_size); @@ -324,27 +408,9 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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; @@ -357,6 +423,27 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigDQKVDtype: std::memcpy(&cfg.dqkv_dtype, buf, attr_size); 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 kNVTEFusedAttnConfigAttnScale: + std::memcpy(&cfg.attn_scale, buf, attr_size); + break; case kNVTEFusedAttnConfigBatchSize: std::memcpy(&cfg.batch_size, buf, attr_size); break; @@ -372,6 +459,27 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigHeadDimV: std::memcpy(&cfg.head_dim_v, 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 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; case kNVTEFusedAttnConfigNumPagesK: std::memcpy(&cfg.num_pages_k, buf, attr_size); break; @@ -402,22 +510,526 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigBiasSeqlenKV: std::memcpy(&cfg.bias_seqlen_kv, buf, attr_size); break; - case kNVTEFusedAttnConfigNumTokensQ: - std::memcpy(&cfg.num_tokens_q, buf, attr_size); + default: + NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + } +} + +NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params() { + + return new transformer_engine::FusedAttnFwdParams( + transformer_engine::make_default_fused_attn_fwd_params()); +} + +void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { + delete transformer_engine::get_fused_attn_fwd_params_mutable(params); +} + +#define NVTE_FWD_PARAMS_GET_BOOL_FIELD(ATTR, FIELD) \ + case ATTR: \ + bool_to_uint8(p.FIELD, buf); \ + break + +#define NVTE_FWD_PARAMS_SET_BOOL_FIELD(ATTR, FIELD) \ + case ATTR: \ + uint8_to_bool(buf, p.FIELD); \ + break + +void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, + "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnFwdParams::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 attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + const auto &p = *get_fused_attn_fwd_params(params); + switch (attr) { + case kNVTEFusedAttnFwdParamsQ: + std::memcpy(buf, &p.Q, attr_size); break; - case kNVTEFusedAttnConfigNumTokensKV: - std::memcpy(&cfg.num_tokens_kv, buf, attr_size); + case kNVTEFusedAttnFwdParamsK: + std::memcpy(buf, &p.K, attr_size); break; - case kNVTEFusedAttnConfigBucketedBatchSize: - std::memcpy(&cfg.bucketed_batch_size, buf, attr_size); + case kNVTEFusedAttnFwdParamsV: + std::memcpy(buf, &p.V, attr_size); break; - case kNVTEFusedAttnConfigBucketedNumTokensQ: - std::memcpy(&cfg.bucketed_num_tokens_q, buf, attr_size); + case kNVTEFusedAttnFwdParamsBias: + std::memcpy(buf, &p.Bias, attr_size); break; - case kNVTEFusedAttnConfigBucketedNumTokensKV: - std::memcpy(&cfg.bucketed_num_tokens_kv, buf, attr_size); + case kNVTEFusedAttnFwdParamsSoftmaxOffset: + std::memcpy(buf, &p.SoftmaxOffset, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQ: + std::memcpy(buf, &p.cu_seqlens_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKV: + std::memcpy(buf, &p.cu_seqlens_kv, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQPadded: + std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKVPadded: + std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableK: + std::memcpy(buf, &p.page_table_k, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableV: + std::memcpy(buf, &p.page_table_v, attr_size); + break; + case kNVTEFusedAttnFwdParamsRngState: + std::memcpy(buf, &p.rng_state, attr_size); + break; + case kNVTEFusedAttnFwdParamsS: + std::memcpy(buf, &p.S, attr_size); + break; + case kNVTEFusedAttnFwdParamsO: + std::memcpy(buf, &p.O, attr_size); + break; + case kNVTEFusedAttnFwdParamsAuxCtxTensors: + std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeLeft: + std::memcpy(buf, &p.window_size_left, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeRight: + std::memcpy(buf, &p.window_size_right, attr_size); + break; + NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsBottomRightDiagonal, + bottom_right_diagonal); + NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsIsTraining, is_training); + NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsReturnMaxLogit, return_max_logit); + NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsCudaGraph, cuda_graph); + case kNVTEFusedAttnFwdParamsWorkspace: + std::memcpy(buf, &p.workspace, attr_size); + break; + case kNVTEFusedAttnFwdParamsStream: + std::memcpy(buf, &p.stream, attr_size); break; default: - NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, + "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; + NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + auto &p = *get_fused_attn_fwd_params_mutable(params); + switch (attr) { + case kNVTEFusedAttnFwdParamsQ: + std::memcpy(&p.Q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsK: + std::memcpy(&p.K, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsV: + std::memcpy(&p.V, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBias: + std::memcpy(&p.Bias, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxOffset: + std::memcpy(&p.SoftmaxOffset, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQ: + std::memcpy(&p.cu_seqlens_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKV: + std::memcpy(&p.cu_seqlens_kv, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQPadded: + std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKVPadded: + std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableK: + std::memcpy(&p.page_table_k, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableV: + std::memcpy(&p.page_table_v, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsRngState: + std::memcpy(&p.rng_state, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsS: + std::memcpy(&p.S, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsO: + std::memcpy(&p.O, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAuxCtxTensors: + std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxType: + std::memcpy(&p.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeLeft: + std::memcpy(&p.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeRight: + std::memcpy(&p.window_size_right, buf, attr_size); + break; + NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsBottomRightDiagonal, + bottom_right_diagonal); + NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsIsTraining, is_training); + NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsReturnMaxLogit, return_max_logit); + NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsCudaGraph, cuda_graph); + case kNVTEFusedAttnFwdParamsWorkspace: + std::memcpy(&p.workspace, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsStream: + std::memcpy(&p.stream, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); } } + +#undef NVTE_FWD_PARAMS_GET_BOOL_FIELD +#undef NVTE_FWD_PARAMS_SET_BOOL_FIELD + +NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params() { + return new transformer_engine::FusedAttnBwdParams( + transformer_engine::make_default_fused_attn_bwd_params()); +} + +void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { + delete transformer_engine::get_fused_attn_bwd_params_mutable(params); +} + +#define NVTE_BWD_PARAMS_GET_BOOL_FIELD(ATTR, FIELD) \ + case ATTR: \ + bool_to_uint8(p.FIELD, buf); \ + break + +#define NVTE_BWD_PARAMS_SET_BOOL_FIELD(ATTR, FIELD) \ + case ATTR: \ + uint8_to_bool(buf, p.FIELD); \ + break + +void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, + "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnBwdParams::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 attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + const auto &p = *get_fused_attn_bwd_params(params); + switch (attr) { + case kNVTEFusedAttnBwdParamsQ: + std::memcpy(buf, &p.Q, attr_size); + break; + case kNVTEFusedAttnBwdParamsK: + std::memcpy(buf, &p.K, attr_size); + break; + case kNVTEFusedAttnBwdParamsV: + std::memcpy(buf, &p.V, attr_size); + break; + case kNVTEFusedAttnBwdParamsO: + std::memcpy(buf, &p.O, attr_size); + break; + case kNVTEFusedAttnBwdParamsDO: + std::memcpy(buf, &p.dO, attr_size); + break; + case kNVTEFusedAttnBwdParamsS: + std::memcpy(buf, &p.S, attr_size); + break; + case kNVTEFusedAttnBwdParamsDP: + std::memcpy(buf, &p.dP, attr_size); + break; + case kNVTEFusedAttnBwdParamsAuxCtxTensors: + std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQ: + std::memcpy(buf, &p.dQ, attr_size); + break; + case kNVTEFusedAttnBwdParamsDK: + std::memcpy(buf, &p.dK, attr_size); + break; + case kNVTEFusedAttnBwdParamsDV: + std::memcpy(buf, &p.dV, attr_size); + break; + case kNVTEFusedAttnBwdParamsDBias: + std::memcpy(buf, &p.dBias, attr_size); + break; + case kNVTEFusedAttnBwdParamsDSoftmaxOffset: + std::memcpy(buf, &p.dSoftmaxOffset, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQ: + std::memcpy(buf, &p.cu_seqlens_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKV: + std::memcpy(buf, &p.cu_seqlens_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: + std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: + std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnBwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOFormat: + std::memcpy(buf, &p.do_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQKVLayout: + std::memcpy(buf, &p.dqkv_layout, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOScaleInvFormat: + std::memcpy(buf, &p.do_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeLeft: + std::memcpy(buf, &p.window_size_left, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeRight: + std::memcpy(buf, &p.window_size_right, attr_size); + break; + NVTE_BWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnBwdParamsBottomRightDiagonal, + bottom_right_diagonal); + NVTE_BWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnBwdParamsDeterministic, deterministic); + NVTE_BWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnBwdParamsCudaGraph, cuda_graph); + case kNVTEFusedAttnBwdParamsWorkspace: + std::memcpy(buf, &p.workspace, attr_size); + break; + case kNVTEFusedAttnBwdParamsStream: + std::memcpy(buf, &p.stream, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, + "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; + NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + auto &p = *get_fused_attn_bwd_params_mutable(params); + switch (attr) { + case kNVTEFusedAttnBwdParamsQ: + std::memcpy(&p.Q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsK: + std::memcpy(&p.K, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsV: + std::memcpy(&p.V, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsO: + std::memcpy(&p.O, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDO: + std::memcpy(&p.dO, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsS: + std::memcpy(&p.S, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDP: + std::memcpy(&p.dP, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAuxCtxTensors: + std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQ: + std::memcpy(&p.dQ, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDK: + std::memcpy(&p.dK, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDV: + std::memcpy(&p.dV, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDBias: + std::memcpy(&p.dBias, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDSoftmaxOffset: + std::memcpy(&p.dSoftmaxOffset, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQ: + std::memcpy(&p.cu_seqlens_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKV: + std::memcpy(&p.cu_seqlens_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: + std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: + std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOFormat: + std::memcpy(&p.do_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQKVLayout: + std::memcpy(&p.dqkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOScaleInvFormat: + std::memcpy(&p.do_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(&p.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeLeft: + std::memcpy(&p.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeRight: + std::memcpy(&p.window_size_right, buf, attr_size); + break; + NVTE_BWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnBwdParamsBottomRightDiagonal, + bottom_right_diagonal); + NVTE_BWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnBwdParamsDeterministic, deterministic); + NVTE_BWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnBwdParamsCudaGraph, cuda_graph); + case kNVTEFusedAttnBwdParamsWorkspace: + std::memcpy(&p.workspace, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsStream: + std::memcpy(&p.stream, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +#undef NVTE_BWD_PARAMS_GET_BOOL_FIELD +#undef NVTE_BWD_PARAMS_SET_BOOL_FIELD diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 025d0166bf..3d38961ead 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -19,122 +19,146 @@ namespace transformer_engine { struct FusedAttnConfig { + // basic attention knobs bool is_training = false; bool deterministic = false; bool cuda_graph = false; bool return_max_logit = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + float dropout = 0.0f; + + // data types + NVTEDType qkv_dtype = kNVTEBFloat16; + NVTEDType o_dtype = kNVTEBFloat16; + NVTEDType do_dtype = kNVTEBFloat16; + NVTEDType dqkv_dtype = kNVTEBFloat16; + + // data and scale layout 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; + + // attention 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; + + // tensor dimensions 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 max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + size_t num_tokens_q = 0; + size_t num_tokens_kv = 0; + + // derived tensor dimensions + size_t bucketed_batch_size = 0; + size_t bucketed_num_tokens_q = 0; + size_t bucketed_num_tokens_kv = 0; + + // paged KV dimensions 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; + + // bias dimensions 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[] = { + // basic attention knobs 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(NVTE_Bias_Type), // bias_type sizeof(int64_t), // window_size_left sizeof(int64_t), // window_size_right sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(NVTEScalingMode), // scaling_mode + sizeof(float), // dropout + // data types sizeof(NVTEDType), // qkv_dtype sizeof(NVTEDType), // o_dtype sizeof(NVTEDType), // do_dtype sizeof(NVTEDType), // dqkv_dtype + // data and scale layout + 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 + // attention scaling + sizeof(float), // attn_scale + // tensor dimensions 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), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(size_t), // num_tokens_q + sizeof(size_t), // num_tokens_kv + // derived tensor dimensions + sizeof(size_t), // bucketed_batch_size + sizeof(size_t), // bucketed_num_tokens_q + sizeof(size_t), // bucketed_num_tokens_kv + // paged KV dimensions 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 + // bias dimensions 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) < + return std::tie(is_training, deterministic, cuda_graph, return_max_logit, attn_mask_type, + bias_type, window_size_left, window_size_right, bottom_right_diagonal, + softmax_type, scaling_mode, dropout, qkv_dtype, o_dtype, do_dtype, dqkv_dtype, + qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, attn_scale, batch_size, num_attn_heads, num_gqa_groups, + head_dim_qk, head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, + num_tokens_kv, bucketed_batch_size, bucketed_num_tokens_q, bucketed_num_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_batch_size, bias_num_heads, bias_seqlen_q, + bias_seqlen_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.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, + rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, + rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, rhs.qkv_layout, + rhs.o_format, rhs.do_format, rhs.dqkv_layout, rhs.qkv_scale_inv_format, + rhs.do_scale_inv_format, rhs.attn_scale, rhs.batch_size, rhs.num_attn_heads, + rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, rhs.max_seqlen_q, + rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, rhs.bucketed_batch_size, + rhs.bucketed_num_tokens_q, rhs.bucketed_num_tokens_kv, 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); + rhs.bias_seqlen_q, rhs.bias_seqlen_kv); } }; @@ -156,6 +180,189 @@ inline FusedAttnConfig *get_fused_attn_config_mutable(NVTEFusedAttnConfig config return reinterpret_cast(config); } +struct FusedAttnFwdParams { + NVTETensor Q = nullptr; + NVTETensor K = nullptr; + NVTETensor V = nullptr; + NVTETensor Bias = nullptr; + NVTETensor SoftmaxOffset = nullptr; + NVTETensor cu_seqlens_q = nullptr; + NVTETensor cu_seqlens_kv = nullptr; + NVTETensor cu_seqlens_q_padded = nullptr; + NVTETensor cu_seqlens_kv_padded = nullptr; + NVTETensor page_table_k = nullptr; + NVTETensor page_table_v = nullptr; + NVTETensor rng_state = nullptr; + NVTETensor S = nullptr; + NVTETensor O = nullptr; + NVTETensorPack *Aux_CTX_Tensors = nullptr; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format qkv_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; + float attn_scale = 1.0f; + float dropout = 0.0f; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + bool is_training = false; + bool return_max_logit = false; + bool cuda_graph = false; + NVTETensor workspace = nullptr; + cudaStream_t stream = nullptr; + + static constexpr size_t attr_sizes[] = { + sizeof(NVTETensor), // Q + sizeof(NVTETensor), // K + sizeof(NVTETensor), // V + sizeof(NVTETensor), // Bias + sizeof(NVTETensor), // SoftmaxOffset + sizeof(NVTETensor), // cu_seqlens_q + sizeof(NVTETensor), // cu_seqlens_kv + sizeof(NVTETensor), // cu_seqlens_q_padded + sizeof(NVTETensor), // cu_seqlens_kv_padded + sizeof(NVTETensor), // page_table_k + sizeof(NVTETensor), // page_table_v + sizeof(NVTETensor), // rng_state + sizeof(NVTETensor), // S + sizeof(NVTETensor), // O + sizeof(NVTETensorPack *), // Aux_CTX_Tensors + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(float), // attn_scale + sizeof(float), // dropout + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(uint8_t), // is_training + sizeof(uint8_t), // return_max_logit + sizeof(uint8_t), // cuda_graph + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream + }; +}; + +struct FusedAttnBwdParams { + NVTETensor Q = nullptr; + NVTETensor K = nullptr; + NVTETensor V = nullptr; + NVTETensor O = nullptr; + NVTETensor dO = nullptr; + NVTETensor S = nullptr; + NVTETensor dP = nullptr; + const NVTETensorPack *Aux_CTX_Tensors = nullptr; + NVTETensor dQ = nullptr; + NVTETensor dK = nullptr; + NVTETensor dV = nullptr; + NVTETensor dBias = nullptr; + NVTETensor dSoftmaxOffset = nullptr; + NVTETensor cu_seqlens_q = nullptr; + NVTETensor cu_seqlens_kv = nullptr; + NVTETensor cu_seqlens_q_padded = nullptr; + NVTETensor cu_seqlens_kv_padded = nullptr; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + 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; + float attn_scale = 1.0f; + float dropout = 0.0f; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + bool deterministic = false; + bool cuda_graph = false; + NVTETensor workspace = nullptr; + cudaStream_t stream = nullptr; + + static constexpr size_t attr_sizes[] = { + sizeof(NVTETensor), // Q + sizeof(NVTETensor), // K + sizeof(NVTETensor), // V + sizeof(NVTETensor), // O + sizeof(NVTETensor), // dO + sizeof(NVTETensor), // S + sizeof(NVTETensor), // dP + sizeof(const NVTETensorPack *), // Aux_CTX_Tensors + sizeof(NVTETensor), // dQ + sizeof(NVTETensor), // dK + sizeof(NVTETensor), // dV + sizeof(NVTETensor), // dBias + sizeof(NVTETensor), // dSoftmaxOffset + sizeof(NVTETensor), // cu_seqlens_q + sizeof(NVTETensor), // cu_seqlens_kv + sizeof(NVTETensor), // cu_seqlens_q_padded + sizeof(NVTETensor), // cu_seqlens_kv_padded + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + 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(float), // attn_scale + sizeof(float), // dropout + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream + }; +}; + +inline FusedAttnFwdParams make_default_fused_attn_fwd_params() { return FusedAttnFwdParams{}; } + +inline FusedAttnBwdParams make_default_fused_attn_bwd_params() { return FusedAttnBwdParams{}; } + +// Build a FusedAttnConfig from the scalar "knobs" carried by the fwd/bwd params (mask/bias/softmax +// type, scales, dropout, window, layout/format fields, flags). The tensor-derived fields (dtypes, +// dims, scaling_mode, paged-KV / bias dims, token counts) are left at their defaults and must be +// filled in by the caller from the actual tensors. +FusedAttnConfig make_fused_attn_config(const FusedAttnFwdParams ¶ms); +FusedAttnConfig make_fused_attn_config(const FusedAttnBwdParams ¶ms); + +inline const FusedAttnFwdParams *get_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); + return reinterpret_cast(params); +} + +inline FusedAttnFwdParams *get_fused_attn_fwd_params_mutable(NVTEFusedAttnFwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); + return reinterpret_cast(params); +} + +inline const FusedAttnBwdParams *get_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); + return reinterpret_cast(params); +} + +inline FusedAttnBwdParams *get_fused_attn_bwd_params_mutable(NVTEFusedAttnBwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); + return reinterpret_cast(params); +} + } // 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 6e1bf518f6..d414dd62a1 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -244,11 +244,10 @@ void set_message(const char **message, std::string reason) { } // namespace // select a backend for fused attention -namespace { - -NVTE_Fused_Attn_Backend select_fused_attn_backend(const transformer_engine::FusedAttnConfig &cfg, - const char **message) { +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, + const char **message) { using namespace transformer_engine; + const FusedAttnConfig &cfg = *get_fused_attn_config(config); set_message(message, ""); cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); @@ -339,14 +338,6 @@ NVTE_Fused_Attn_Backend select_fused_attn_backend(const transformer_engine::Fuse 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, @@ -381,44 +372,32 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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); + return nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), + /*message=*/nullptr); } -// NVTE fused attention FWD with separate Q, K and V -void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - const NVTETensor page_table_v, const NVTETensor rng_state, - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, - bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_flash_attn_fwd); +void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { + NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(cu_seqlens_kv_padded); - const Tensor *input_page_table_k = convertNVTETensorCheck(page_table_k); - const Tensor *input_page_table_v = convertNVTETensorCheck(page_table_v); - const Tensor *input_rng_state = convertNVTETensorCheck(rng_state); - const Tensor *input_Q = convertNVTETensorCheck(Q); - const Tensor *input_K = convertNVTETensorCheck(K); - const Tensor *input_V = convertNVTETensorCheck(V); - const Tensor *input_Bias = convertNVTETensorCheck(Bias); - const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(SoftmaxOffset); - Tensor *input_output_S = convertNVTETensorCheck(S); - Tensor *output_O = convertNVTETensorCheck(O); - Tensor *wkspace = convertNVTETensor(workspace); - - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const FusedAttnFwdParams &p = *get_fused_attn_fwd_params(params); + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); + const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(p.cu_seqlens_q_padded); + const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(p.cu_seqlens_kv_padded); + const Tensor *input_page_table_k = convertNVTETensorCheck(p.page_table_k); + const Tensor *input_page_table_v = convertNVTETensorCheck(p.page_table_v); + const Tensor *input_rng_state = convertNVTETensorCheck(p.rng_state); + const Tensor *input_Q = convertNVTETensorCheck(p.Q); + const Tensor *input_K = convertNVTETensorCheck(p.K); + const Tensor *input_V = convertNVTETensorCheck(p.V); + const Tensor *input_Bias = convertNVTETensorCheck(p.Bias); + const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(p.SoftmaxOffset); + Tensor *input_output_S = convertNVTETensorCheck(p.S); + Tensor *output_O = convertNVTETensorCheck(p.O); + Tensor *wkspace = convertNVTETensor(p.workspace); + + NVTE_QKV_Format q_format = nvte_get_q_format(p.qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(p.qkv_layout); auto *q_dims = input_Q->data.shape.data(); auto *k_dims = input_K->data.shape.data(); auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING @@ -447,15 +426,15 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso if (input_page_table_v->data.dptr != nullptr) { max_pages_per_seq_v = input_page_table_v->data.shape[1]; } - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(p.qkv_layout); if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - if (kv_format == NVTE_QKV_Format::NVTE_BSHD) { + NVTE_QKV_Format paged_kv_format = nvte_get_kv_format(p.qkv_layout); + if (paged_kv_format == NVTE_QKV_Format::NVTE_BSHD) { num_pages_k = input_K->data.shape[0]; page_size_k = input_K->data.shape[1]; num_pages_v = input_V->data.shape[0]; page_size_v = input_V->data.shape[1]; - } else if (kv_format == NVTE_QKV_Format::NVTE_SBHD) { + } else if (paged_kv_format == NVTE_QKV_Format::NVTE_SBHD) { num_pages_k = input_K->data.shape[1]; page_size_k = input_K->data.shape[0]; num_pages_v = input_V->data.shape[1]; @@ -471,7 +450,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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) && + if ((p.bias_type != NVTE_NO_BIAS) && (p.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]; @@ -479,25 +458,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso bias_skv = input_Bias->data.shape[3]; } - 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; + FusedAttnConfig cfg = make_fused_attn_config(p); 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; @@ -518,62 +480,103 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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); + nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), + /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - cfg.is_training = is_training; + cfg.is_training = p.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, + output_O, p.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); + input_page_table_v, input_rng_state, wkspace, p.stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - 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); + cfg.is_training = p.is_training; + fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, + p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, + wkspace, p.stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } } -// NVTE fused attention BWD with separate Q, K and V -void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, - const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, - NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, + +// NVTE fused attention FWD with separate Q, K and V +void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, + const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, + NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, - size_t max_seqlen_kv, float attn_scale, float dropout, + const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, + const NVTETensor page_table_v, const NVTETensor rng_state, + size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, + bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, - bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_flash_attn_bwd); + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, + NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_flash_attn_fwd); + transformer_engine::FusedAttnFwdParams p = transformer_engine::make_default_fused_attn_fwd_params(); + p.Q = Q; + p.K = K; + p.V = V; + p.Bias = Bias; + p.SoftmaxOffset = SoftmaxOffset; + p.S = S; + p.O = O; + p.Aux_CTX_Tensors = Aux_CTX_Tensors; + p.cu_seqlens_q = cu_seqlens_q; + p.cu_seqlens_kv = cu_seqlens_kv; + p.cu_seqlens_q_padded = cu_seqlens_q_padded; + p.cu_seqlens_kv_padded = cu_seqlens_kv_padded; + p.page_table_k = page_table_k; + p.page_table_v = page_table_v; + p.rng_state = rng_state; + p.max_seqlen_q = max_seqlen_q; + p.max_seqlen_kv = max_seqlen_kv; + p.is_training = is_training; + p.return_max_logit = return_max_logit; + p.cuda_graph = cuda_graph; + p.attn_scale = attn_scale; + p.dropout = dropout; + p.qkv_layout = qkv_layout; + p.o_format = o_format; + p.qkv_scale_inv_format = qkv_scale_inv_format; + p.bias_type = bias_type; + p.attn_mask_type = attn_mask_type; + p.softmax_type = softmax_type; + p.window_size_left = window_size_left; + p.window_size_right = window_size_right; + p.bottom_right_diagonal = bottom_right_diagonal; + p.workspace = workspace; + p.stream = stream; + nvte_fused_attn_fwd_v2(reinterpret_cast(&p)); +} + +void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { + NVTE_API_CALL(nvte_fused_attn_bwd_v2); using namespace transformer_engine; - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(cu_seqlens_kv_padded); - const Tensor *input_Q = convertNVTETensorCheck(Q); - const Tensor *input_K = convertNVTETensorCheck(K); - const Tensor *input_V = convertNVTETensorCheck(V); - const Tensor *input_O = convertNVTETensorCheck(O); - const Tensor *input_dO = convertNVTETensorCheck(dO); - const Tensor *input_S = convertNVTETensorCheck(S); - Tensor *input_output_dP = convertNVTETensorCheck(dP); - Tensor *output_dQ = convertNVTETensorCheck(dQ); - Tensor *output_dK = convertNVTETensorCheck(dK); - Tensor *output_dV = convertNVTETensorCheck(dV); - Tensor *output_dBias = convertNVTETensorCheck(dBias); - Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(dSoftmaxOffset); - Tensor *wkspace = convertNVTETensor(workspace); - - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const FusedAttnBwdParams &p = *get_fused_attn_bwd_params(params); + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); + const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(p.cu_seqlens_q_padded); + const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(p.cu_seqlens_kv_padded); + const Tensor *input_Q = convertNVTETensorCheck(p.Q); + const Tensor *input_K = convertNVTETensorCheck(p.K); + const Tensor *input_V = convertNVTETensorCheck(p.V); + const Tensor *input_O = convertNVTETensorCheck(p.O); + const Tensor *input_dO = convertNVTETensorCheck(p.dO); + const Tensor *input_S = convertNVTETensorCheck(p.S); + Tensor *input_output_dP = convertNVTETensorCheck(p.dP); + Tensor *output_dQ = convertNVTETensorCheck(p.dQ); + Tensor *output_dK = convertNVTETensorCheck(p.dK); + Tensor *output_dV = convertNVTETensorCheck(p.dV); + Tensor *output_dBias = convertNVTETensorCheck(p.dBias); + Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(p.dSoftmaxOffset); + Tensor *wkspace = convertNVTETensor(p.workspace); + + NVTE_QKV_Format q_format = nvte_get_q_format(p.qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(p.qkv_layout); auto *q_dims = input_Q->data.shape.data(); auto *k_dims = input_K->data.shape.data(); auto *v_dims = input_V->data.shape.data(); @@ -598,7 +601,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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) && + if ((p.bias_type != NVTE_NO_BIAS) && (p.bias_type != NVTE_ALIBI) && output_dBias->data.shape.size() >= 4) { bias_b = output_dBias->data.shape[0]; bias_h = output_dBias->data.shape[1]; @@ -606,28 +609,8 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso bias_skv = output_dBias->data.shape[3]; } - 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; + FusedAttnConfig cfg = make_fused_attn_config(p); 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; @@ -644,46 +627,105 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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); + nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), + /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + Tensor *output_S = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); Tensor *input_Bias, *input_SoftmaxOffset; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + if ((p.bias_type != NVTE_NO_BIAS) && (p.bias_type != NVTE_ALIBI)) { + input_Bias = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); } - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { + input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); } 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, + input_cu_seqlens_kv_padded, input_rng_state, wkspace, p.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++]); - const Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + const Tensor *input_M = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + const Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); const Tensor *input_SoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { + input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); } const Tensor *input_dO_f16 = nullptr; if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { - input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + input_dO_f16 = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); } 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); + input_rng_state, wkspace, p.stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } } +// NVTE fused attention BWD with separate Q, K and V +void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, + const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, + const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, + NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, + const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, + size_t max_seqlen_kv, float attn_scale, float dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, + bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_flash_attn_bwd); + transformer_engine::FusedAttnBwdParams p = transformer_engine::make_default_fused_attn_bwd_params(); + p.Q = Q; + p.K = K; + p.V = V; + p.O = O; + p.dO = dO; + p.S = S; + p.dP = dP; + p.Aux_CTX_Tensors = Aux_CTX_Tensors; + p.dQ = dQ; + p.dK = dK; + p.dV = dV; + p.dBias = dBias; + p.dSoftmaxOffset = dSoftmaxOffset; + p.cu_seqlens_q = cu_seqlens_q; + p.cu_seqlens_kv = cu_seqlens_kv; + p.cu_seqlens_q_padded = cu_seqlens_q_padded; + p.cu_seqlens_kv_padded = cu_seqlens_kv_padded; + p.max_seqlen_q = max_seqlen_q; + p.max_seqlen_kv = max_seqlen_kv; + p.attn_scale = attn_scale; + p.dropout = dropout; + p.qkv_layout = qkv_layout; + p.o_format = o_format; + p.do_format = do_format; + p.dqkv_layout = dqkv_layout; + p.qkv_scale_inv_format = qkv_scale_inv_format; + p.do_scale_inv_format = do_scale_inv_format; + p.bias_type = bias_type; + p.attn_mask_type = attn_mask_type; + p.softmax_type = softmax_type; + p.window_size_left = window_size_left; + p.window_size_right = window_size_right; + p.bottom_right_diagonal = bottom_right_diagonal; + p.deterministic = deterministic; + p.cuda_graph = cuda_graph; + p.workspace = workspace; + p.stream = stream; + nvte_fused_attn_bwd_v2(reinterpret_cast(&p)); +} + uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlen, NVTETensor workspace, size_t len, cudaStream_t stream) { NVTE_API_CALL(nvte_get_runtime_num_segments); 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 11bccc09f4..f96e212787 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 @@ -256,8 +256,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( sdpa_options = fe::graph::SDPA_attributes() .set_name("flash_attention") .set_generate_stats(generate_stats) - .set_causal_mask(is_causal) - .set_causal_mask_bottom_right(is_bottom_right) .set_attn_scale(attn_scale); fe::DiagonalAlignment_t const &diagonal_alignment = @@ -269,6 +267,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_options.set_diagonal_band_right_bound(window_size_right); + } else if (is_causal || is_bottom_right) { + // Preferred replacement for the deprecated set_causal_mask[_bottom_right]: causal + // masking = diagonal alignment (set above) + a right band bound of 0. + sdpa_options.set_diagonal_band_right_bound(0); } sdpa_options.set_alibi_mask(is_alibi); @@ -607,12 +609,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const int sm_arch_ = cuda::sm_arch(device_id); bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); - if (is_paged_kv) { - NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); - } - // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { @@ -773,8 +769,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( fe::graph::SDPA_backward_attributes sdpa_backward_options; sdpa_backward_options = fe::graph::SDPA_backward_attributes() .set_name("flash_attention_backward") - .set_causal_mask(is_causal) - .set_causal_mask_bottom_right(is_bottom_right) .set_attn_scale(attn_scale); if (use_ragged_stats) { @@ -794,6 +788,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } else if (is_causal || is_bottom_right) { + // Preferred replacement for the deprecated set_causal_mask[_bottom_right]: causal + // masking = diagonal alignment (set above) + a right band bound of 0. + sdpa_backward_options.set_diagonal_band_right_bound(0); } if (cudnn_runtime_version >= 90000) { @@ -1077,16 +1075,8 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrS1 = nullptr; void *devPtrS2 = nullptr; void *devPtrBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - size_t bias_sq = 0; - size_t bias_skv = 0; if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { devPtrBias = input_Bias->data.dptr; - 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]; } void *devPtrSoftmaxOffset = nullptr; if (softmax_type != NVTE_VANILLA_SOFTMAX) { @@ -1105,12 +1095,6 @@ void fused_attn_arbitrary_seqlen_fwd( 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; if (Aux_CTX_Tensors->size == 0) { @@ -1145,7 +1129,8 @@ void fused_attn_arbitrary_seqlen_fwd( if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_bias->data.dptr = nullptr; - output_bias->data.shape = {bias_b, bias_h, bias_sq, bias_skv}; + output_bias->data.shape = {graph_cfg.bias_batch_size, graph_cfg.bias_num_heads, + graph_cfg.bias_seqlen_q, graph_cfg.bias_seqlen_kv}; output_bias->data.dtype = QKV_type; } @@ -1226,27 +1211,13 @@ void fused_attn_arbitrary_seqlen_bwd( void *devPtrdO = input_dO->data.dptr; void *devPtrBias = nullptr; void *devPtrdBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - size_t bias_sq = 0; - size_t bias_skv = 0; if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { devPtrBias = input_Bias->data.dptr; devPtrdBias = output_dBias->data.dptr; - 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]; } 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; void *devPtrdK = output_dK->data.dptr; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 90d24d2b36..e40d606648 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -218,7 +218,6 @@ void fused_attn_fp8_fwd_impl( sdpa_options = fe::graph::SDPA_fp8_attributes() .set_name("sdpa_fp8") .set_generate_stats(true) - .set_causal_mask(is_causal) .set_attn_scale(attn_scale); fe::DiagonalAlignment_t const& diagonal_alignment = @@ -234,6 +233,12 @@ void fused_attn_fp8_fwd_impl( sdpa_options.set_diagonal_band_right_bound(window_size_right); } } + // Preferred replacement for the deprecated set_causal_mask: causal masking = diagonal + // alignment (set above) + a right band bound of 0, unless an explicit right bound was + // already applied above. + if (is_causal && !(cudnn_runtime_version >= 92100 && window_size_right != -1)) { + sdpa_options.set_diagonal_band_right_bound(0); + } // sdpa_options.set_alibi_mask(is_alibi); // if (is_bias) { @@ -765,7 +770,6 @@ void fused_attn_fp8_bwd_impl( fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; sdpa_backward_options = fe::graph::SDPA_fp8_backward_attributes() .set_name("sdpa_fp8_backward") - .set_causal_mask(is_causal) .set_attn_scale(attn_scale); fe::DiagonalAlignment_t const& diagonal_alignment = @@ -781,6 +785,12 @@ void fused_attn_fp8_bwd_impl( sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); } } + // Preferred replacement for the deprecated set_causal_mask: causal masking = diagonal + // alignment (set above) + a right band bound of 0, unless an explicit right bound was + // already applied above. + if (is_causal && !(cudnn_runtime_version >= 92100 && window_size_right != -1)) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } // sdpa_backward_options.set_alibi_mask(is_alibi); diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index c338f1a99d..44413b40ef 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -636,43 +636,6 @@ __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/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index f26e03c5ad..685f069488 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -202,55 +202,64 @@ 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. + * This enum is used to index the ``FusedAttnConfig`` struct. The order of its fields must match that of + * the declaration fields and the ``attr_sizes`` array of that struct. New fields may only be appended + * at the end, and existing fields are never to be reordered, removed, or resized. */ enum NVTEFusedAttnConfigAttribute { + // basic attention knobs kNVTEFusedAttnConfigIsTraining = 0, kNVTEFusedAttnConfigDeterministic, kNVTEFusedAttnConfigCudaGraph, kNVTEFusedAttnConfigReturnMaxLogit, - kNVTEFusedAttnConfigQKVLayout, - kNVTEFusedAttnConfigOFormat, - kNVTEFusedAttnConfigDOFormat, - kNVTEFusedAttnConfigDQKVLayout, - kNVTEFusedAttnConfigQKVScaleInvFormat, - kNVTEFusedAttnConfigDOScaleInvFormat, - kNVTEFusedAttnConfigBiasType, kNVTEFusedAttnConfigAttnMaskType, - kNVTEFusedAttnConfigSoftmaxType, - kNVTEFusedAttnConfigScalingMode, - kNVTEFusedAttnConfigAttnScale, - kNVTEFusedAttnConfigDropout, - kNVTEFusedAttnConfigMaxSeqlenQ, - kNVTEFusedAttnConfigMaxSeqlenKV, + kNVTEFusedAttnConfigBiasType, kNVTEFusedAttnConfigWindowSizeLeft, kNVTEFusedAttnConfigWindowSizeRight, kNVTEFusedAttnConfigBottomRightDiagonal, + kNVTEFusedAttnConfigSoftmaxType, + kNVTEFusedAttnConfigScalingMode, + kNVTEFusedAttnConfigDropout, + // data types kNVTEFusedAttnConfigQKVDtype, kNVTEFusedAttnConfigODtype, kNVTEFusedAttnConfigDODtype, kNVTEFusedAttnConfigDQKVDtype, + // data and scale layout + kNVTEFusedAttnConfigQKVLayout, + kNVTEFusedAttnConfigOFormat, + kNVTEFusedAttnConfigDOFormat, + kNVTEFusedAttnConfigDQKVLayout, + kNVTEFusedAttnConfigQKVScaleInvFormat, + kNVTEFusedAttnConfigDOScaleInvFormat, + // attention scaling + kNVTEFusedAttnConfigAttnScale, + // tensor dimensions kNVTEFusedAttnConfigBatchSize, kNVTEFusedAttnConfigNumAttnHeads, kNVTEFusedAttnConfigNumGqaGroups, kNVTEFusedAttnConfigHeadDimQK, kNVTEFusedAttnConfigHeadDimV, + kNVTEFusedAttnConfigMaxSeqlenQ, + kNVTEFusedAttnConfigMaxSeqlenKV, + kNVTEFusedAttnConfigNumTokensQ, + kNVTEFusedAttnConfigNumTokensKV, + // derived tensor dimensions + kNVTEFusedAttnConfigBucketedBatchSize, + kNVTEFusedAttnConfigBucketedNumTokensQ, + kNVTEFusedAttnConfigBucketedNumTokensKV, + // paged KV dimensions kNVTEFusedAttnConfigNumPagesK, kNVTEFusedAttnConfigNumPagesV, kNVTEFusedAttnConfigPageSizeK, kNVTEFusedAttnConfigPageSizeV, kNVTEFusedAttnConfigMaxPagesPerSeqK, kNVTEFusedAttnConfigMaxPagesPerSeqV, + // bias dimensions kNVTEFusedAttnConfigBiasBatchSize, kNVTEFusedAttnConfigBiasNumHeads, kNVTEFusedAttnConfigBiasSeqlenQ, kNVTEFusedAttnConfigBiasSeqlenKV, - kNVTEFusedAttnConfigNumTokensQ, - kNVTEFusedAttnConfigNumTokensKV, - kNVTEFusedAttnConfigBucketedBatchSize, - kNVTEFusedAttnConfigBucketedNumTokensQ, - kNVTEFusedAttnConfigBucketedNumTokensKV, kNVTEFusedAttnConfigNumAttributes }; @@ -279,6 +288,128 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, NVTEFusedAttnConfigAttribute attr, const void *buf, size_t size_in_bytes); +/*! \brief Opaque fused-attention forward-parameter handle. */ +typedef void *NVTEFusedAttnFwdParams; + +/*! \enum NVTEFusedAttnFwdParamsAttribute + * \brief Attribute types for ``NVTEFusedAttnFwdParams``. + */ +enum NVTEFusedAttnFwdParamsAttribute { + kNVTEFusedAttnFwdParamsQ = 0, + kNVTEFusedAttnFwdParamsK, + kNVTEFusedAttnFwdParamsV, + kNVTEFusedAttnFwdParamsBias, + kNVTEFusedAttnFwdParamsSoftmaxOffset, + kNVTEFusedAttnFwdParamsCuSeqlensQ, + kNVTEFusedAttnFwdParamsCuSeqlensKV, + kNVTEFusedAttnFwdParamsCuSeqlensQPadded, + kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, + kNVTEFusedAttnFwdParamsPageTableK, + kNVTEFusedAttnFwdParamsPageTableV, + kNVTEFusedAttnFwdParamsRngState, + kNVTEFusedAttnFwdParamsS, + kNVTEFusedAttnFwdParamsO, + kNVTEFusedAttnFwdParamsAuxCtxTensors, + kNVTEFusedAttnFwdParamsMaxSeqlenQ, + kNVTEFusedAttnFwdParamsMaxSeqlenKV, + kNVTEFusedAttnFwdParamsQKVLayout, + kNVTEFusedAttnFwdParamsOFormat, + kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + kNVTEFusedAttnFwdParamsBiasType, + kNVTEFusedAttnFwdParamsAttnMaskType, + kNVTEFusedAttnFwdParamsSoftmaxType, + kNVTEFusedAttnFwdParamsAttnScale, + kNVTEFusedAttnFwdParamsDropout, + kNVTEFusedAttnFwdParamsWindowSizeLeft, + kNVTEFusedAttnFwdParamsWindowSizeRight, + kNVTEFusedAttnFwdParamsBottomRightDiagonal, + kNVTEFusedAttnFwdParamsIsTraining, + kNVTEFusedAttnFwdParamsReturnMaxLogit, + kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsWorkspace, + kNVTEFusedAttnFwdParamsStream, + kNVTEFusedAttnFwdParamsNumAttributes +}; + +/*! \brief Create a default-initialized fused-attention forward-parameter object. */ +NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void); + +/*! \brief Destroy a fused-attention forward-parameter handle. */ +void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); + +/*! \brief Query an attribute in a fused-attention forward-parameter object. */ +void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention forward-parameter object. */ +void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, const void *buf, + size_t size_in_bytes); + +/*! \brief Opaque fused-attention backward-parameter handle. */ +typedef void *NVTEFusedAttnBwdParams; + +/*! \enum NVTEFusedAttnBwdParamsAttribute + * \brief Attribute types for ``NVTEFusedAttnBwdParams``. + */ +enum NVTEFusedAttnBwdParamsAttribute { + kNVTEFusedAttnBwdParamsQ = 0, + kNVTEFusedAttnBwdParamsK, + kNVTEFusedAttnBwdParamsV, + kNVTEFusedAttnBwdParamsO, + kNVTEFusedAttnBwdParamsDO, + kNVTEFusedAttnBwdParamsS, + kNVTEFusedAttnBwdParamsDP, + kNVTEFusedAttnBwdParamsAuxCtxTensors, + kNVTEFusedAttnBwdParamsDQ, + kNVTEFusedAttnBwdParamsDK, + kNVTEFusedAttnBwdParamsDV, + kNVTEFusedAttnBwdParamsDBias, + kNVTEFusedAttnBwdParamsDSoftmaxOffset, + kNVTEFusedAttnBwdParamsCuSeqlensQ, + kNVTEFusedAttnBwdParamsCuSeqlensKV, + kNVTEFusedAttnBwdParamsCuSeqlensQPadded, + kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, + kNVTEFusedAttnBwdParamsMaxSeqlenQ, + kNVTEFusedAttnBwdParamsMaxSeqlenKV, + kNVTEFusedAttnBwdParamsQKVLayout, + kNVTEFusedAttnBwdParamsOFormat, + kNVTEFusedAttnBwdParamsDOFormat, + kNVTEFusedAttnBwdParamsDQKVLayout, + kNVTEFusedAttnBwdParamsQKVScaleInvFormat, + kNVTEFusedAttnBwdParamsDOScaleInvFormat, + kNVTEFusedAttnBwdParamsBiasType, + kNVTEFusedAttnBwdParamsAttnMaskType, + kNVTEFusedAttnBwdParamsSoftmaxType, + kNVTEFusedAttnBwdParamsAttnScale, + kNVTEFusedAttnBwdParamsDropout, + kNVTEFusedAttnBwdParamsWindowSizeLeft, + kNVTEFusedAttnBwdParamsWindowSizeRight, + kNVTEFusedAttnBwdParamsBottomRightDiagonal, + kNVTEFusedAttnBwdParamsDeterministic, + kNVTEFusedAttnBwdParamsCudaGraph, + kNVTEFusedAttnBwdParamsWorkspace, + kNVTEFusedAttnBwdParamsStream, + kNVTEFusedAttnBwdParamsNumAttributes +}; + +/*! \brief Create a default-initialized fused-attention backward-parameter object. */ +NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void); + +/*! \brief Destroy a fused-attention backward-parameter handle. */ +void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params); + +/*! \brief Query an attribute in a fused-attention backward-parameter object. */ +void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention backward-parameter object. */ +void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, const void *buf, + size_t size_in_bytes); + /*! \brief Get fused attention backend based on input parameters. * * This call exercises cudnn-frontend's support checks by building (and caching) @@ -393,6 +524,12 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ +void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); + +/*! \brief Compute dot product attention with separate Q, K and V. + * + * \deprecated This function has been deprecated in favor of nvte_fused_attn_fwd_v2. + */ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, @@ -466,6 +603,12 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ +void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); + +/*! \brief Compute the backward of the dot product attention with separate Q, K and V. + * + * \deprecated This function has been deprecated in favor of nvte_fused_attn_bwd_v2. + */ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, @@ -1059,6 +1202,352 @@ class FusedAttnConfigWrapper { NVTEFusedAttnConfig cfg_ = nullptr; }; +/*! \class FusedAttnFwdParamsWrapper + * \brief C++ helper for constructing an ``NVTEFusedAttnFwdParams``. + */ +class FusedAttnFwdParamsWrapper { + public: + FusedAttnFwdParamsWrapper() : params_{nvte_create_fused_attn_fwd_params()} {} + FusedAttnFwdParamsWrapper(const FusedAttnFwdParamsWrapper &) = delete; + FusedAttnFwdParamsWrapper &operator=(const FusedAttnFwdParamsWrapper &) = delete; + FusedAttnFwdParamsWrapper(FusedAttnFwdParamsWrapper &&other) noexcept : params_{other.params_} { + other.params_ = nullptr; + } + FusedAttnFwdParamsWrapper &operator=(FusedAttnFwdParamsWrapper &&other) noexcept { + if (this != &other) { + nvte_destroy_fused_attn_fwd_params(params_); + params_ = other.params_; + other.params_ = nullptr; + } + return *this; + } + ~FusedAttnFwdParamsWrapper() { + if (params_ != nullptr) { + nvte_destroy_fused_attn_fwd_params(params_); + } + } + operator NVTEFusedAttnFwdParams() const noexcept { return params_; } + NVTEFusedAttnFwdParams get() const noexcept { return params_; } + FusedAttnFwdParamsWrapper &set_Q(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQ, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_K(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsK, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_V(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsV, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_Bias(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBias, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_SoftmaxOffset(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxOffset, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQ, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKV, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQPadded, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_page_table_k(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableK, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_page_table_v(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableV, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_rng_state(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsRngState, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsS, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsO, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack * val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAuxCtxTensors, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenQ, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenKV, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVLayout, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsOFormat, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVScaleInvFormat, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBiasType, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnMaskType, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxType, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnScale, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsDropout, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeLeft, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeRight, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBottomRightDiagonal, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsIsTraining, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsReturnMaxLogit, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCudaGraph, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWorkspace, &val, sizeof(val)); + return *this; + } + FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsStream, &val, sizeof(val)); + return *this; + } + private: + NVTEFusedAttnFwdParams params_ = nullptr; +}; + +/*! \class FusedAttnBwdParamsWrapper + * \brief C++ helper for constructing an ``NVTEFusedAttnBwdParams``. + */ +class FusedAttnBwdParamsWrapper { + public: + FusedAttnBwdParamsWrapper() : params_{nvte_create_fused_attn_bwd_params()} {} + FusedAttnBwdParamsWrapper(const FusedAttnBwdParamsWrapper &) = delete; + FusedAttnBwdParamsWrapper &operator=(const FusedAttnBwdParamsWrapper &) = delete; + FusedAttnBwdParamsWrapper(FusedAttnBwdParamsWrapper &&other) noexcept : params_{other.params_} { + other.params_ = nullptr; + } + FusedAttnBwdParamsWrapper &operator=(FusedAttnBwdParamsWrapper &&other) noexcept { + if (this != &other) { + nvte_destroy_fused_attn_bwd_params(params_); + params_ = other.params_; + other.params_ = nullptr; + } + return *this; + } + ~FusedAttnBwdParamsWrapper() { + if (params_ != nullptr) { + nvte_destroy_fused_attn_bwd_params(params_); + } + } + operator NVTEFusedAttnBwdParams() const noexcept { return params_; } + NVTEFusedAttnBwdParams get() const noexcept { return params_; } + FusedAttnBwdParamsWrapper &set_Q(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQ, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_K(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsK, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_V(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsV, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_O(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsO, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dO(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDO, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_S(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsS, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dP(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDP, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(const NVTETensorPack * val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAuxCtxTensors, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dQ(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQ, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dK(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDK, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dV(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDV, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dBias(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDBias, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dSoftmaxOffset(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDSoftmaxOffset, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQ, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKV, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQPadded, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenQ, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenKV, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVLayout, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsOFormat, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOFormat, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQKVLayout, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOScaleInvFormat, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBiasType, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnMaskType, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsSoftmaxType, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnScale, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDropout, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeLeft, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeRight, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBottomRightDiagonal, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDeterministic, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCudaGraph, &u8_val, sizeof(u8_val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWorkspace, &val, sizeof(val)); + return *this; + } + FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsStream, &val, sizeof(val)); + return *this; + } + private: + NVTEFusedAttnBwdParams params_ = nullptr; +}; #endif // __cplusplus #endif 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 f3078c1a9c..d5b743129c 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 @@ -1537,32 +1537,11 @@ def forward( _alibi_cache["_alibi_slopes_require_update"] = True _alibi_cache["_alibi_bias_require_update"] = True - # detect bias shape - core_attention_bias_shape = None - if core_attention_bias is not None: - if ( - core_attention_bias.shape[0] == batch_size - and core_attention_bias.shape[1] == query_layer.shape[-2] - ): - core_attention_bias_shape = "bhss" - elif ( - core_attention_bias.shape[0] == 1 - and core_attention_bias.shape[1] == query_layer.shape[-2] - ): - core_attention_bias_shape = "1hss" - elif ( - core_attention_bias.shape[0] == batch_size and core_attention_bias.shape[1] == 1 - ): - core_attention_bias_shape = "b1ss" - elif core_attention_bias.shape[0] == 1 and core_attention_bias.shape[1] == 1: - if core_attention_bias.shape[2] == 1: - core_attention_bias_shape = "111s" - else: - core_attention_bias_shape = "11ss" - else: - assert ( - False - ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss, 111s} shapes" + core_attention_bias_shape = ( + tuple(core_attention_bias.shape) + if core_attention_bias_type != "no_bias" and core_attention_bias is not None + else None + ) # Default pad_between_seqs auto-detect. For THD, infer presence of # inter-sequence padding from whether padded cu_seqlens were supplied -- @@ -1629,16 +1608,20 @@ def forward( num_gqa_groups=num_gqa_groups, max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv, + num_tokens_q=(query_layer.shape[0] if q_format == "thd" else 0), + num_tokens_kv=(key_layer.shape[0] if kv_format == "thd" else 0), head_dim_qk=head_dim_qk, head_dim_v=head_dim_v, attn_mask_type=attn_mask_type, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, - alibi_slopes_shape=alibi_slopes.shape if alibi_slopes is not None else None, + alibi_slopes_shape=alibi_slopes.shape if core_attention_bias_type == "alibi" and alibi_slopes is not None else None, core_attention_bias_type=core_attention_bias_type, core_attention_bias_shape=core_attention_bias_shape, core_attention_bias_requires_grad=( - core_attention_bias.requires_grad if core_attention_bias is not None else False + core_attention_bias.requires_grad + if core_attention_bias_type != "no_bias" and core_attention_bias is not None + else False ), pad_between_seqs=pad_between_seqs, attention_dropout=self.attention_dropout, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 68bb0b199e..b981a809a5 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -47,7 +47,7 @@ from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.pytorch.quantization import get_fp8_te_dtype -from transformer_engine.pytorch.constants import TE_DType, MXFP8_BLOCK_SCALING_SIZE +from transformer_engine.pytorch.constants import TE_DType, DType, MXFP8_BLOCK_SCALING_SIZE from transformer_engine.pytorch.utils import ( @@ -211,6 +211,10 @@ class AttentionParams: Maximum sequence length of the query tensor. max_seqlen_kv : int, default = 128 Maximum sequence length of the key and value tensors. + num_tokens_q : int, default = 0 + Total number of query tokens in a batch, when `qkv_format=thd`. + num_tokens_kv : int, default = 0 + Total number of key/value tokens in a batch, when `qkv_format=thd`. head_dim_qk : int, default = 64 The size of each attention head in query and key tensors. head_dim_v : int, default = 64 @@ -227,8 +231,10 @@ class AttentionParams: Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. - core_attention_bias_shape : str, default = 1hss - Attention bias shape, {`1hss`, `b1ss`, `bhss`}. + core_attention_bias_shape : Optional[Tuple[int, int, int, int]], default = None + Broadcast shape of the `core_attention_bias` tensor as `(b, h, sq, skv)`. `None` when no + bias tensor is present. The broadcast pattern (`1hss`, `bhss`, etc.) is derived inside + `get_attention_backend`. core_attention_bias_requires_grad : bool, default = True Whether attention bias requires gradient. pad_between_seqs : bool, default = False @@ -281,6 +287,8 @@ class AttentionParams: num_gqa_groups: int = 16 max_seqlen_q: int = 128 max_seqlen_kv: int = 128 + num_tokens_q: int = 0 + num_tokens_kv: int = 0 head_dim_qk: int = 64 head_dim_v: int = 64 attn_mask_type: str = "no_mask" @@ -288,7 +296,7 @@ class AttentionParams: bottom_right_diagonal: bool = True alibi_slopes_shape: Union[torch.Size, List, None] = None core_attention_bias_type: str = "no_bias" - core_attention_bias_shape: str = "1hss" + core_attention_bias_shape: Union[Tuple[int, int, int, int], None] = None core_attention_bias_requires_grad: bool = True pad_between_seqs: bool = False attention_dropout: float = 0.0 @@ -329,6 +337,74 @@ def __eq__(self, other): return True +@dataclass(eq=True) +class FusedAttentionParams: + """ + Attention parameters used by the `FusedAttention` backend. + """ + + # basic attention knobs + is_training: bool = True + deterministic: bool = False + cuda_graph: bool = False + return_max_logit: bool = False + attn_mask_type: tex.NVTE_Mask_Type = tex.NVTE_Mask_Type.NVTE_NO_MASK + bias_type: tex.NVTE_Bias_Type = tex.NVTE_Bias_Type.NVTE_NO_BIAS + window_size_left: int = -1 + window_size_right: int = -1 + bottom_right_diagonal: bool = True + softmax_type: tex.NVTE_Softmax_Type = tex.NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX + scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_INVALID_SCALING + dropout: float = 0.0 + + # data types + qkv_dtype: DType = DType.kBFloat16 + o_dtype: DType = DType.kBFloat16 + do_dtype: DType = DType.kBFloat16 + dqkv_dtype: DType = DType.kBFloat16 + + # data and scale layout + qkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + o_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + do_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + dqkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + qkv_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + do_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + + # attention scaling + attn_scale: float = 0.0 + + # tensor dimensions + batch_size: int = 0 + num_attn_heads: int = 0 + num_gqa_groups: int = 0 + head_dim_qk: int = 0 + head_dim_v: int = 0 + max_seqlen_q: int = 0 + max_seqlen_kv: int = 0 + num_tokens_q: int = 0 + num_tokens_kv: int = 0 + + # derived tensor dimensions + bucketed_batch_size: int = 0 + bucketed_num_tokens_q: int = 0 + bucketed_num_tokens_kv: int = 0 + + # paged KV dimensions + num_pages_k: int = 0 + num_pages_v: int = 0 + page_size_k: int = 0 + page_size_v: int = 0 + max_pages_per_seq_k: int = 0 + max_pages_per_seq_v: int = 0 + + # bias dimensions + bias_batch_size: int = 0 + bias_num_heads: int = 0 + bias_seqlen_q: int = 0 + bias_seqlen_kv: int = 0 + + def get_attention_backend( attention_params: AttentionParams = None, ): @@ -364,6 +440,8 @@ def get_attention_backend( num_gqa_groups = attention_params.num_gqa_groups max_seqlen_q = attention_params.max_seqlen_q max_seqlen_kv = attention_params.max_seqlen_kv + num_tokens_q = attention_params.num_tokens_q + num_tokens_kv = attention_params.num_tokens_kv head_dim_qk = attention_params.head_dim_qk head_dim_v = attention_params.head_dim_v attn_mask_type = attention_params.attn_mask_type @@ -1340,42 +1418,57 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fu_core_attention_bias_requires_grad = False if len(alibi_slopes_shape) == 1 and alibi_slopes_shape[0] == num_heads: - fu_core_attention_bias_shape = "1hss" + fu_core_attention_bias_shape = (1, num_heads, max_seqlen_q, max_seqlen_kv) elif ( len(alibi_slopes_shape) == 2 and alibi_slopes_shape[0] == batch_size and alibi_slopes_shape[1] == num_heads ): - fu_core_attention_bias_shape = "bhss" - + fu_core_attention_bias_shape = (batch_size, num_heads, max_seqlen_q, max_seqlen_kv) + + fu_core_attention_bias_shape_type = None + if fu_core_attention_bias_type == "post_scale_bias" and fu_core_attention_bias_shape is not None: + b, h, sq, _skv = fu_core_attention_bias_shape + if b == batch_size and h == num_heads: + fu_core_attention_bias_shape_type = "bhss" + elif b == 1 and h == num_heads: + fu_core_attention_bias_shape_type = "1hss" + elif b == batch_size and h == 1: + fu_core_attention_bias_shape_type = "b1ss" + elif b == 1 and h == 1: + fu_core_attention_bias_shape_type = "111s" if sq == 1 and max_seqlen_q != 1 else "11ss" + else: + raise ValueError( + f"core_attention_bias tensor must be in one of " + "{"bhss", "1hss", "b1ss", "11ss", "111s"} shapes. Found (b,h,sq,skv) = ({b},{h},{sq},{_skv})" + ) if ( use_fused_attention and fu_core_attention_bias_type == "post_scale_bias" - and fu_core_attention_bias_shape != "1hss" + and fu_core_attention_bias_shape_type != "1hss" ): # dbias calculation is not supported for 111s as of cuDNN 9.18. So, use fused attention backend only if bias does not require grad. - if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape == "111s": + if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape_type == "111s": logger.warning( "Disabling FusedAttention as dbias calculation is not supported for 111s" ) use_fused_attention = False + # Filter: cuDNN support fused_attention_backend = None if use_fused_attention: # ``DType`` is implicitly convertible to ``transformer_engine::DType`` # 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 + qkv_type = TE_DType[qkv_dtype] + o_type = qkv_type + do_type = qkv_type + dqkv_type = qkv_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: recipe = fp8_meta["recipe"] - q_type = get_fp8_te_dtype(recipe, fprop_tensor=True) - kv_type = q_type + qkv_type = get_fp8_te_dtype(recipe, fprop_tensor=True) 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 @@ -1391,45 +1484,67 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt dqkv_type = TE_DType[torch.bfloat16] else: scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING - o_type = q_type + o_type = qkv_type do_type = o_type - dqkv_type = q_type + dqkv_type = qkv_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, - max_seqlen_q, - max_seqlen_kv, - head_dim_qk, - head_dim_v, - window_size[0], - window_size[1], - bottom_right_diagonal, - return_max_logit, - cuda_graph, - deterministic, + num_pages_k = num_pages_v = 0 + page_size_k = page_size_v = 0 + max_pages_per_seq_k = max_pages_per_seq_v = 0 + if inference_params is not None and getattr(inference_params, "is_paged", False): + num_pages_k = num_pages_v = inference_params.total_num_pages + page_size_k = page_size_v = inference_params.page_size + max_pages_per_seq_k = max_pages_per_seq_v = inference_params.cache_manager.max_pages_per_seq + bias_batch_size = bias_num_heads = bias_seqlen_q = bias_seqlen_kv = 0 + if fu_core_attention_bias_shape is not None: + bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = fu_core_attention_bias_shape + fused_attn_params = FusedAttentionParams( + is_training=is_training, + deterministic=deterministic, + cuda_graph=cuda_graph, + return_max_logit=return_max_logit, + attn_mask_type=AttnMaskType[attn_mask_type], + bias_type=AttnBiasType[fu_core_attention_bias_type], + window_size_left=window_size[0], + window_size_right=window_size[1], + bottom_right_diagonal=bottom_right_diagonal, + softmax_type=SoftmaxType[softmax_type], + scaling_mode=scaling_mode, + dropout=attention_dropout, + qkv_dtype=qkv_type, + o_dtype=o_type, + do_dtype=do_type, + dqkv_dtype=dqkv_type, + qkv_layout=QKVLayout[qkv_layout], + o_format=QKVFormat[o_format], + do_format=QKVFormat[do_format], + dqkv_layout=QKVLayout[dqkv_layout], + qkv_scale_inv_format=QKVFormat[qkv_scale_inv_format], + do_scale_inv_format=QKVFormat[do_scale_inv_format], + attn_scale=softmax_scale, + batch_size=batch_size, + num_attn_heads=num_heads, + num_gqa_groups=num_gqa_groups, + head_dim_qk=head_dim_qk, + head_dim_v=head_dim_v, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, + num_pages_k=num_pages_k, + num_pages_v=num_pages_v, + page_size_k=page_size_k, + page_size_v=page_size_v, + max_pages_per_seq_k=max_pages_per_seq_k, + max_pages_per_seq_v=max_pages_per_seq_v, + bias_batch_size=bias_batch_size, + bias_num_heads=bias_num_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ) + fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug( "Disabling FusedAttention: %s", diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 982fbdb169..84c831f23e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -89,15 +89,7 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T // 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); + py::object fused_attn_params); 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 706eb630e0..464a409063 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -41,49 +41,53 @@ namespace transformer_engine::pytorch { // get the 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."); - + py::object fused_attn_params) { + py::object &p = fused_attn_params; 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); - + cfg.set_is_training(p.attr("is_training").cast()) + .set_deterministic(p.attr("deterministic").cast()) + .set_cuda_graph(p.attr("cuda_graph").cast()) + .set_return_max_logit(p.attr("return_max_logit").cast()) + .set_attn_mask_type(p.attr("attn_mask_type").cast()) + .set_bias_type(p.attr("bias_type").cast()) + .set_window_size_left(p.attr("window_size_left").cast()) + .set_window_size_right(p.attr("window_size_right").cast()) + .set_bottom_right_diagonal(p.attr("bottom_right_diagonal").cast()) + .set_softmax_type(p.attr("softmax_type").cast()) + .set_scaling_mode(p.attr("scaling_mode").cast()) + .set_dropout(p.attr("dropout").cast()) + .set_qkv_dtype(static_cast(p.attr("qkv_dtype").cast())) + .set_o_dtype(static_cast(p.attr("o_dtype").cast())) + .set_do_dtype(static_cast(p.attr("do_dtype").cast())) + .set_dqkv_dtype(static_cast(p.attr("dqkv_dtype").cast())) + .set_qkv_layout(p.attr("qkv_layout").cast()) + .set_o_format(p.attr("o_format").cast()) + .set_do_format(p.attr("do_format").cast()) + .set_dqkv_layout(p.attr("dqkv_layout").cast()) + .set_qkv_scale_inv_format(p.attr("qkv_scale_inv_format").cast()) + .set_do_scale_inv_format(p.attr("do_scale_inv_format").cast()) + .set_attn_scale(p.attr("attn_scale").cast()) + .set_batch_size(p.attr("batch_size").cast()) + .set_num_attn_heads(p.attr("num_attn_heads").cast()) + .set_num_gqa_groups(p.attr("num_gqa_groups").cast()) + .set_head_dim_qk(p.attr("head_dim_qk").cast()) + .set_head_dim_v(p.attr("head_dim_v").cast()) + .set_max_seqlen_q(p.attr("max_seqlen_q").cast()) + .set_max_seqlen_kv(p.attr("max_seqlen_kv").cast()) + .set_num_tokens_q(p.attr("num_tokens_q").cast()) + .set_num_tokens_kv(p.attr("num_tokens_kv").cast()) + .set_num_pages_k(p.attr("num_pages_k").cast()) + .set_num_pages_v(p.attr("num_pages_v").cast()) + .set_page_size_k(p.attr("page_size_k").cast()) + .set_page_size_v(p.attr("page_size_v").cast()) + .set_max_pages_per_seq_k(p.attr("max_pages_per_seq_k").cast()) + .set_max_pages_per_seq_v(p.attr("max_pages_per_seq_v").cast()) + .set_bias_batch_size(p.attr("bias_batch_size").cast()) + .set_bias_num_heads(p.attr("bias_num_heads").cast()) + .set_bias_seqlen_q(p.attr("bias_seqlen_q").cast()) + .set_bias_seqlen_kv(p.attr("bias_seqlen_kv").cast()); + + py::gil_scoped_release nogil; 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()}; diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 9c9ec36138..d9050ab941 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -411,7 +411,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Swap first two tensor dimensions", py::arg("tensor"), py::kw_only(), py::arg("out"), py::call_guard()); m.def("get_fused_attn_backend", &transformer_engine::pytorch::get_fused_attn_backend, - "Get Fused Attention backend", py::call_guard()); + "Get Fused Attention backend", py::arg("fused_attn_params")); m.def("compute_amax", &transformer_engine::pytorch::compute_amax, "Compute absolute max value in tensor", py::arg("input"), py::arg("amax"), py::call_guard()); From 88a327c1dc33c3ac425c71492911953f67bdf704 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:27:18 -0700 Subject: [PATCH 26/63] remove bucketed b/t_q/t_kv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.cpp | 20 +------------------ .../common/fused_attn/config_and_params.h | 8 ++------ .../include/transformer_engine/fused_attn.h | 19 ------------------ .../attention/dot_product_attention/utils.py | 5 ----- 4 files changed, 3 insertions(+), 49 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 79976fab14..0aee8fce55 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -135,7 +135,7 @@ FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg) { // cuDNN graph supports dynamic shapes for batch_size cache_cfg.batch_size = 1; cache_cfg.bucketed_batch_size = 1; - cache_cfg.attention_scale = 1.0f; + cache_cfg.attn_scale = 1.0f; return cache_cfg; } @@ -315,15 +315,6 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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; case kNVTEFusedAttnConfigNumPagesK: std::memcpy(buf, &cfg.num_pages_k, attr_size); break; @@ -471,15 +462,6 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, 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; case kNVTEFusedAttnConfigNumPagesK: std::memcpy(&cfg.num_pages_k, buf, attr_size); break; diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 3d38961ead..6ae69a6197 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -48,7 +48,7 @@ struct FusedAttnConfig { NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; // attention scaling - float attn_scale = 0.0f; + float attn_scale = 1.0f; // tensor dimensions size_t batch_size = 0; @@ -61,7 +61,7 @@ struct FusedAttnConfig { size_t num_tokens_q = 0; size_t num_tokens_kv = 0; - // derived tensor dimensions + // derived tensor dimensions (internal only) size_t bucketed_batch_size = 0; size_t bucketed_num_tokens_q = 0; size_t bucketed_num_tokens_kv = 0; @@ -118,10 +118,6 @@ struct FusedAttnConfig { sizeof(size_t), // max_seqlen_kv sizeof(size_t), // num_tokens_q sizeof(size_t), // num_tokens_kv - // derived tensor dimensions - sizeof(size_t), // bucketed_batch_size - sizeof(size_t), // bucketed_num_tokens_q - sizeof(size_t), // bucketed_num_tokens_kv // paged KV dimensions sizeof(size_t), // num_pages_k sizeof(size_t), // num_pages_v diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 685f069488..36601d2b2c 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -244,10 +244,6 @@ enum NVTEFusedAttnConfigAttribute { kNVTEFusedAttnConfigMaxSeqlenKV, kNVTEFusedAttnConfigNumTokensQ, kNVTEFusedAttnConfigNumTokensKV, - // derived tensor dimensions - kNVTEFusedAttnConfigBucketedBatchSize, - kNVTEFusedAttnConfigBucketedNumTokensQ, - kNVTEFusedAttnConfigBucketedNumTokensKV, // paged KV dimensions kNVTEFusedAttnConfigNumPagesK, kNVTEFusedAttnConfigNumPagesV, @@ -1182,21 +1178,6 @@ class FusedAttnConfigWrapper { 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; diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index b981a809a5..6ea0237848 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -385,11 +385,6 @@ class FusedAttentionParams: num_tokens_q: int = 0 num_tokens_kv: int = 0 - # derived tensor dimensions - bucketed_batch_size: int = 0 - bucketed_num_tokens_q: int = 0 - bucketed_num_tokens_kv: int = 0 - # paged KV dimensions num_pages_k: int = 0 num_pages_v: int = 0 From 391fe2eb4776fec97bc5ddd17956234d21128dc2 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 14 Jul 2026 14:33:18 -0700 Subject: [PATCH 27/63] thread _v2 through, fix default scaling mode, make config specific to fwd/bwd, thread bias shapes through in jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.cpp | 37 ++-- .../common/fused_attn/config_and_params.h | 8 +- .../common/fused_attn/fused_attn.cpp | 12 +- .../fused_attn_f16_arbitrary_seqlen.cu | 4 +- .../common/fused_attn/fused_attn_fp8.cu | 4 +- .../jax/cpp_extensions/attention.py | 44 ++++ transformer_engine/jax/csrc/extensions.h | 3 +- .../jax/csrc/extensions/attention.cpp | 206 ++++++++++++++---- .../attention/dot_product_attention/utils.py | 4 +- .../pytorch/csrc/extensions/attention.cpp | 116 +++++++--- 10 files changed, 322 insertions(+), 116 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 0aee8fce55..0dcae6da5f 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -37,7 +37,6 @@ 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); @@ -45,7 +44,6 @@ void populate_fused_attn_config(FusedAttnConfig *cfg) { 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); @@ -81,24 +79,9 @@ void populate_fused_attn_config(FusedAttnConfig *cfg) { 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); - } - } } -FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg) { +FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg, bool is_forward) { FusedAttnConfig cache_cfg = cfg; const int64_t s_q = static_cast(cache_cfg.max_seqlen_q); @@ -137,12 +120,28 @@ FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg) { cache_cfg.bucketed_batch_size = 1; cache_cfg.attn_scale = 1.0f; + // Drop from each graph's cache key the fields the graph does not actually consume, so a graph + // prewarmed by a backend probe (which may carry different values for those ignored fields, e.g. + // the framework get_attention_backend probe) is still reused at execution. The forward graph + // produces O (and optionally softmax stats / max logit) but never consumes the dO/dQKV dtypes or + // the backward-only determinism choice. The backward graph consumes dO/dQKV and honors + // determinism but never produces the forward max-logit output. + if (is_forward) { + if (cache_cfg.is_training) { + cache_cfg.do_dtype = kNVTEBFloat16; + cache_cfg.dqkv_dtype = kNVTEBFloat16; + cache_cfg.deterministic = false; + } + } else { + cache_cfg.return_max_logit = false; + } + return cache_cfg; } FusedAttnConfig make_fused_attn_config(const FusedAttnFwdParams ¶ms) { FusedAttnConfig cfg = make_default_fused_attn_config(); - cfg.is_training = false; // fwd-only probe; caller restores before dispatch + cfg.is_training = params.is_training; cfg.deterministic = false; cfg.cuda_graph = params.cuda_graph; cfg.return_max_logit = params.return_max_logit; diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 6ae69a6197..f7d0c684aa 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -163,8 +163,12 @@ 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); +// bottom-right mask folding). Call after populate_fused_attn_config(). Pass is_forward=true when +// keying a forward graph and is_forward=false for a backward graph; each key drops the fields the +// corresponding graph does not consume so it is not fragmented by them: a training forward key +// drops the dO/dQKV dtypes and the (backward-only) deterministic flag, and a backward key drops +// the (forward-only) return_max_logit flag. +FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg, bool is_forward); inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index d414dd62a1..5737475bc9 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -348,13 +348,12 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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.attn_scale = attn_scale; cfg.dropout = dropout; cfg.max_seqlen_q = max_seqlen_q; cfg.max_seqlen_kv = max_seqlen_kv; @@ -363,13 +362,14 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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.o_dtype = q_dtype; + cfg.do_dtype = q_dtype; + cfg.dqkv_dtype = q_dtype; 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.is_training = is_training; cfg.return_max_logit = return_max_logit; cfg.deterministic = deterministic; return nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), @@ -484,14 +484,12 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { /*message=*/nullptr); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - cfg.is_training = p.is_training; fused_attn_arbitrary_seqlen_fwd(cfg, input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset, output_O, p.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, p.stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - cfg.is_training = p.is_training; fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, p.stream, handle); 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 f96e212787..c4918991d0 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 @@ -142,7 +142,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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); + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg, /*is_forward=*/true); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -625,7 +625,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( // 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); + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg, /*is_forward=*/false); try { namespace fe = cudnn_frontend; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index e40d606648..8da0ad2261 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -81,7 +81,7 @@ 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); + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg, /*is_forward=*/true); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -514,7 +514,7 @@ 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); + const FusedAttnConfig cache_cfg = make_fused_attn_graph_cache_config(cfg, /*is_forward=*/false); try { namespace fe = cudnn_frontend; using graph_and_tensors = diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index eaa9c8769a..0d328b734d 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -130,6 +130,13 @@ class FusedAttnHelper: window_size: Tuple[int, int] bottom_right_diagonal: bool attn_scale: float = 1.0 + # Actual POST_SCALE_BIAS operand dims (may be broadcast, e.g. 1). Left None when the caller does + # not know the bias shape (e.g. the config-level is_fused_attn_kernel_available API), in which + # case get_fused_attn_backend falls back to the full [b, h, sq, skv] representative shape. + bias_batch: Optional[int] = None + bias_heads: Optional[int] = None + bias_seqlen_q: Optional[int] = None + bias_seqlen_kv: Optional[int] = None def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel. @@ -147,6 +154,22 @@ def get_fused_attn_backend(self): diagnostic string describing why the configuration was rejected when backend = NVTE_No_Backend. """ q_type = jax_dtype_to_te_dtype(self.q_dtype) + # The support probe builds a cuDNN graph, which for POST_SCALE_BIAS needs a concrete bias + # shape. Prefer the actual bias operand dims (threaded from the bias aval) so the probe keys + # and builds the exact graph execution uses, even for broadcast bias. When the caller does + # not know the shape (e.g. the config-level is_fused_attn_kernel_available API), fall back to + # the full [b, h, sq, skv] representative shape; backend support does not depend on the bias + # broadcast pattern. For other bias types there is no bias operand, so pass 0 to avoid + # fragmenting the graph-cache key. + if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + bias_batch = self.bias_batch if self.bias_batch is not None else self.batch_size + bias_heads = self.bias_heads if self.bias_heads is not None else self.q_num_heads + bias_seqlen_q = self.bias_seqlen_q if self.bias_seqlen_q is not None else self.q_max_seqlen + bias_seqlen_kv = ( + self.bias_seqlen_kv if self.bias_seqlen_kv is not None else self.kv_max_seqlen + ) + else: + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 return transformer_engine_jax.get_fused_attn_backend( self.is_training, self.batch_size, @@ -177,6 +200,10 @@ def get_fused_attn_backend(self): self.window_size[1], self.bottom_right_diagonal, not self.is_non_deterministic_allowed(), + bias_batch, + bias_heads, + bias_seqlen_q, + bias_seqlen_kv, ) @staticmethod @@ -366,6 +393,19 @@ def abstract( # backend determines the softmax buffer shape/dtype input_batch = reduce(operator.mul, batch_shape) + # Thread the actual POST_SCALE_BIAS operand dims so the trace-time support probe keys and + # builds the exact cuDNN graph the runtime executes (incl. broadcast bias). Derive them the + # same way the lowering/execution does: the bias aval is [*batch, heads, sq, skv] where the + # leading batch dims may be >1 and are collapsed into a single bias_batch. Matching that + # computation keeps the prewarm and runtime graph-cache keys identical for any bias rank. + is_post_scale_bias = config.attn_bias_type == AttnBiasType.POST_SCALE_BIAS + if is_post_scale_bias: + *probe_bias_batch_shape, probe_bias_heads, probe_bias_seqlen_q, probe_bias_seqlen_kv = ( + bias_aval.shape + ) + probe_bias_batch = reduce(operator.mul, probe_bias_batch_shape) + else: + probe_bias_batch = probe_bias_heads = probe_bias_seqlen_q = probe_bias_seqlen_kv = None backend, message = FusedAttnHelper( config.is_training, input_batch, @@ -385,6 +425,10 @@ def abstract( config.window_size, config.bottom_right_diagonal, attn_scale=float(config.scaling_factor), + bias_batch=probe_bias_batch, + bias_heads=probe_bias_heads, + bias_seqlen_q=probe_bias_seqlen_q, + bias_seqlen_kv=probe_bias_seqlen_kv, ).get_fused_attn_backend() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 5a6790793c..7e66b35f5b 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -161,7 +161,8 @@ std::tuple GetFusedAttnBackend( 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); + bool deterministic, size_t bias_batch, size_t bias_heads, size_t bias_seqlen_q, + size_t bias_seqlen_kv); 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 c88f63a5e9..d95b1db69f 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -33,7 +33,8 @@ std::tuple GetFusedAttnBackend( 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) { + bool deterministic, size_t bias_batch, size_t bias_heads, size_t bias_seqlen_q, + size_t bias_seqlen_kv) { if (o_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { o_format = nvte_get_q_format(qkv_layout); } @@ -75,7 +76,11 @@ std::tuple GetFusedAttnBackend( .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); + .set_head_dim_v(v_head_dim) + .set_bias_batch_size(bias_batch) + .set_bias_num_heads(bias_heads) + .set_bias_seqlen_q(bias_seqlen_q) + .set_bias_seqlen_kv(bias_seqlen_kv); const char *message = nullptr; auto backend = nvte_get_fused_attn_backend_v2(cfg, &message); @@ -239,15 +244,41 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); auto ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - nvte_fused_attn_fwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, - q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), - ragged_offset_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), nullptr); + FusedAttnFwdParamsWrapper params; + params.set_Q(q_tensor.data()) + .set_K(k_tensor.data()) + .set_V(v_tensor.data()) + .set_Bias(bias_tensor.data()) + .set_SoftmaxOffset(dummy_softmax_offset_tensor.data()) + .set_S(s_tensor.data()) + .set_O(o_tensor.data()) + .set_Aux_CTX_Tensors(&aux_output_tensors) + .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) + .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) + .set_cu_seqlens_q_padded(ragged_offset_tensor.data()) + .set_cu_seqlens_kv_padded(ragged_offset_tensor.data()) + .set_page_table_k(dummy_page_table_tensor.data()) + .set_page_table_v(dummy_page_table_tensor.data()) + .set_rng_state(dummy_rng_state_tensor.data()) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_is_training(is_training) + .set_return_max_logit(false) + .set_cuda_graph(false) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .set_qkv_layout(qkv_layout) + .set_o_format(nvte_get_q_format(qkv_layout)) + .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) + .set_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_workspace(query_workspace_tensor.data()) + .set_stream(nullptr); + nvte_fused_attn_fwd_v2(params); } nvte_tensor_pack_destroy(&aux_output_tensors); @@ -325,7 +356,8 @@ static void FusedAttnForwardImpl( 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); + v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, deterministic, + bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); 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) */ @@ -386,15 +418,41 @@ static void FusedAttnForwardImpl( auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); - nvte_fused_attn_fwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, - q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(), - k_seq_offsets_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); + FusedAttnFwdParamsWrapper params; + params.set_Q(q_tensor.data()) + .set_K(k_tensor.data()) + .set_V(v_tensor.data()) + .set_Bias(bias_tensor.data()) + .set_SoftmaxOffset(softmax_offset_tensor.data()) + .set_S(s_tensor.data()) + .set_O(o_tensor.data()) + .set_Aux_CTX_Tensors(&aux_output_tensors) + .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) + .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) + .set_cu_seqlens_q_padded(q_seq_offsets_tensor.data()) + .set_cu_seqlens_kv_padded(k_seq_offsets_tensor.data()) + .set_page_table_k(dummy_page_table_tensor.data()) + .set_page_table_v(dummy_page_table_tensor.data()) + .set_rng_state(rng_state_tensor.data()) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_is_training(is_training) + .set_return_max_logit(false) + .set_cuda_graph(false) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .set_qkv_layout(qkv_layout) + .set_o_format(nvte_get_q_format(qkv_layout)) + .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) + .set_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_workspace(workspace_tensor.data()) + .set_stream(stream); + nvte_fused_attn_fwd_v2(params); nvte_tensor_pack_destroy(&aux_output_tensors); } @@ -542,19 +600,45 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( auto dummy_ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - nvte_fused_attn_bwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), - doutput_tensor.data(), - s_tensor.data(), // not used for F16 - s_tensor.data(), // not used for F16 - &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), - dbias_tensor.data(), dummy_d_softmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), - kv_cu_seqlens_tensor.data(), dummy_ragged_offset_tensor.data(), - dummy_ragged_offset_tensor.data(), q_max_seqlen, kv_max_seqlen, scaling_factor, - dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - nvte_get_q_format(qkv_layout), qkv_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, false, query_workspace_tensor.data(), nullptr); + FusedAttnBwdParamsWrapper params; + params.set_Q(q_tensor.data()) + .set_K(k_tensor.data()) + .set_V(v_tensor.data()) + .set_O(output_tensor.data()) + .set_dO(doutput_tensor.data()) + .set_S(s_tensor.data()) // not used for F16 + .set_dP(s_tensor.data()) // not used for F16 + .set_Aux_CTX_Tensors(&aux_input_tensors) + .set_dQ(dq_tensor.data()) + .set_dK(dk_tensor.data()) + .set_dV(dv_tensor.data()) + .set_dBias(dbias_tensor.data()) + .set_dSoftmaxOffset(dummy_d_softmax_offset_tensor.data()) + .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) + .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) + .set_cu_seqlens_q_padded(dummy_ragged_offset_tensor.data()) + .set_cu_seqlens_kv_padded(dummy_ragged_offset_tensor.data()) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .set_qkv_layout(qkv_layout) + .set_o_format(nvte_get_q_format(qkv_layout)) + .set_do_format(nvte_get_q_format(qkv_layout)) + .set_dqkv_layout(qkv_layout) + .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) + .set_do_scale_inv_format(NVTE_QKV_Format_NOT_SET) + .set_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_deterministic(deterministic) + .set_cuda_graph(false) + .set_workspace(query_workspace_tensor.data()) + .set_stream(nullptr); + nvte_fused_attn_bwd_v2(params); } nvte_tensor_pack_destroy(&aux_input_tensors); @@ -604,7 +688,8 @@ static void FusedAttnBackwardImpl( 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); + v_head_dim, window_size_left, window_size_right, bottom_right_diagonal, deterministic, + bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); 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); @@ -681,18 +766,45 @@ static void FusedAttnBackwardImpl( } } - nvte_fused_attn_bwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), - doutput_tensor.data(), - s_tensor.data(), // not used for F16 - s_tensor.data(), // not used for F16 - &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), dbias_tensor.data(), - dsoftmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), - q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), q_max_seqlen, kv_max_seqlen, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - nvte_get_q_format(qkv_layout), qkv_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, false, workspace_tensor.data(), stream); + FusedAttnBwdParamsWrapper params; + params.set_Q(q_tensor.data()) + .set_K(k_tensor.data()) + .set_V(v_tensor.data()) + .set_O(output_tensor.data()) + .set_dO(doutput_tensor.data()) + .set_S(s_tensor.data()) // not used for F16 + .set_dP(s_tensor.data()) // not used for F16 + .set_Aux_CTX_Tensors(&aux_input_tensors) + .set_dQ(dq_tensor.data()) + .set_dK(dk_tensor.data()) + .set_dV(dv_tensor.data()) + .set_dBias(dbias_tensor.data()) + .set_dSoftmaxOffset(dsoftmax_offset_tensor.data()) + .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) + .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) + .set_cu_seqlens_q_padded(q_seq_offsets_tensor.data()) + .set_cu_seqlens_kv_padded(k_seq_offsets_tensor.data()) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .set_qkv_layout(qkv_layout) + .set_o_format(nvte_get_q_format(qkv_layout)) + .set_do_format(nvte_get_q_format(qkv_layout)) + .set_dqkv_layout(qkv_layout) + .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) + .set_do_scale_inv_format(NVTE_QKV_Format_NOT_SET) + .set_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_deterministic(deterministic) + .set_cuda_graph(false) + .set_workspace(workspace_tensor.data()) + .set_stream(stream); + nvte_fused_attn_bwd_v2(params); nvte_tensor_pack_destroy(&aux_input_tensors); } diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6ea0237848..dea269a2a8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -354,7 +354,7 @@ class FusedAttentionParams: window_size_right: int = -1 bottom_right_diagonal: bool = True softmax_type: tex.NVTE_Softmax_Type = tex.NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX - scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_INVALID_SCALING + scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING dropout: float = 0.0 # data types @@ -1458,7 +1458,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt o_type = qkv_type do_type = qkv_type dqkv_type = qkv_type - scaling_mode = tex.NVTEScalingMode.NVTE_INVALID_SCALING + scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING qkv_scale_inv_format = None do_scale_inv_format = None if fp8 and fp8_meta["recipe"].fp8_dpa: diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 464a409063..0514f49587 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -282,16 +282,45 @@ std::vector fused_attn_fwd( // create workspace TensorWrapper workspace; + // build the parameter object + FusedAttnFwdParamsWrapper params; + params.set_Q(te_Q.data()) + .set_K(te_K.data()) + .set_V(te_V.data()) + .set_Bias(te_Bias.data()) + .set_SoftmaxOffset(te_SoftmaxOffset.data()) + .set_S(te_S.data()) + .set_O(te_O.data()) + .set_Aux_CTX_Tensors(&nvte_aux_tensor_pack) + .set_cu_seqlens_q(te_cu_seqlens_q.data()) + .set_cu_seqlens_kv(te_cu_seqlens_kv.data()) + .set_cu_seqlens_q_padded(te_cu_seqlens_q_padded.data()) + .set_cu_seqlens_kv_padded(te_cu_seqlens_kv_padded.data()) + .set_page_table_k(te_page_table_k.data()) + .set_page_table_v(te_page_table_v.data()) + .set_rng_state(te_rng_state.data()) + .set_max_seqlen_q(max_seqlen_q) + .set_max_seqlen_kv(max_seqlen_kv) + .set_is_training(is_training) + .set_return_max_logit(return_max_logit) + .set_cuda_graph(cuda_graph) + .set_attn_scale(attn_scale) + .set_dropout(p_dropout) + .set_qkv_layout(qkv_layout) + .set_o_format(o_format) + .set_qkv_scale_inv_format(qkv_scale_inv_format) + .set_bias_type(bias_type) + .set_attn_mask_type(attn_mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size[0]) + .set_window_size_right(window_size[1]) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_stream(at::cuda::getCurrentCUDAStream()); + // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( - te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_SoftmaxOffset.data(), te_S.data(), - te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), - te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, - return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, - qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); + params.set_workspace(workspace.data()); + nvte_fused_attn_fwd_v2(params); }); // allocate memory for workspace and auxiliary output tensors @@ -341,14 +370,8 @@ std::vector fused_attn_fwd( // execute the kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( - te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_SoftmaxOffset.data(), te_S.data(), - te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), - te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, - return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, - qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); + params.set_workspace(workspace.data()); + nvte_fused_attn_fwd_v2(params); }); // destroy tensor wrappers, but not allocated memory @@ -610,17 +633,49 @@ std::vector fused_attn_bwd( // create workspace TensorWrapper workspace; + // build the parameter object + FusedAttnBwdParamsWrapper params; + params.set_Q(te_Q.data()) + .set_K(te_K.data()) + .set_V(te_V.data()) + .set_O(te_O.data()) + .set_dO(te_dO.data()) + .set_S(te_S.data()) + .set_dP(te_dP.data()) + .set_Aux_CTX_Tensors(&nvte_aux_tensor_pack) + .set_dQ(te_dQ.data()) + .set_dK(te_dK.data()) + .set_dV(te_dV.data()) + .set_dBias(te_dBias.data()) + .set_dSoftmaxOffset(te_dSoftmaxOffset.data()) + .set_cu_seqlens_q(te_cu_seqlens_q.data()) + .set_cu_seqlens_kv(te_cu_seqlens_kv.data()) + .set_cu_seqlens_q_padded(te_cu_seqlens_q_padded.data()) + .set_cu_seqlens_kv_padded(te_cu_seqlens_kv_padded.data()) + .set_max_seqlen_q(max_seqlen_q) + .set_max_seqlen_kv(max_seqlen_kv) + .set_attn_scale(attn_scale) + .set_dropout(p_dropout) + .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_window_size_left(window_size[0]) + .set_window_size_right(window_size[1]) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_deterministic(deterministic) + .set_cuda_graph(cuda_graph) + .set_stream(at::cuda::getCurrentCUDAStream()); + // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( - te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), - &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), - te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, - attn_scale, p_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[0], - window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), - at::cuda::getCurrentCUDAStream()); + params.set_workspace(workspace.data()); + nvte_fused_attn_bwd_v2(params); }); // allocate memory for workspace @@ -630,15 +685,8 @@ std::vector fused_attn_bwd( // execute kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( - te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), - &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), - te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, - attn_scale, p_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[0], - window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), - at::cuda::getCurrentCUDAStream()); + params.set_workspace(workspace.data()); + nvte_fused_attn_bwd_v2(params); }); // destroy tensor wrappers From 261bb9a8eee868676c6bdeb7d73209ae810c5df8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 21:35:35 +0000 Subject: [PATCH 28/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/config_and_params.cpp | 17 +- .../common/fused_attn/config_and_params.h | 162 ++++++++-------- .../common/fused_attn/fused_attn.cpp | 27 +-- .../fused_attn_f16_arbitrary_seqlen.cu | 34 ++-- .../fused_attn_f16_arbitrary_seqlen.h | 36 ++-- .../common/fused_attn/fused_attn_fp8.cu | 58 +++--- .../common/fused_attn/fused_attn_fp8.h | 26 +-- .../include/transformer_engine/fused_attn.h | 178 ++++++++++++------ .../jax/cpp_extensions/attention.py | 4 +- .../jax/csrc/extensions/attention.cpp | 26 ++- .../dot_product_attention.py | 6 +- 11 files changed, 325 insertions(+), 249 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 0dcae6da5f..1c7bb6ae4e 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -202,8 +202,8 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, size_t size_in_bytes, size_t *size_written) { using namespace transformer_engine; - NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, - "Invalid NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + 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; @@ -213,8 +213,8 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, } 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)"); + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, + " bytes)"); const auto &cfg = *get_fused_attn_config(config); switch (attr) { @@ -354,13 +354,13 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, size_t size_in_bytes) { using namespace transformer_engine; - NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, - "Invalid NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + 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)"); + 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); @@ -497,7 +497,6 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, } NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params() { - return new transformer_engine::FusedAttnFwdParams( transformer_engine::make_default_fused_attn_fwd_params()); } diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index f7d0c684aa..5f31ecb0ed 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -11,11 +11,11 @@ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ +#include + #include "common/common.h" #include "transformer_engine/fused_attn.h" -#include - namespace transformer_engine { struct FusedAttnConfig { @@ -82,54 +82,54 @@ struct FusedAttnConfig { static constexpr size_t attr_sizes[] = { // basic attention knobs - sizeof(uint8_t), // is_training - sizeof(uint8_t), // deterministic - sizeof(uint8_t), // cuda_graph - sizeof(uint8_t), // return_max_logit - sizeof(NVTE_Mask_Type), // attn_mask_type - sizeof(NVTE_Bias_Type), // bias_type - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(NVTEScalingMode), // scaling_mode - sizeof(float), // dropout + sizeof(uint8_t), // is_training + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(NVTEScalingMode), // scaling_mode + sizeof(float), // dropout // data types - sizeof(NVTEDType), // qkv_dtype - sizeof(NVTEDType), // o_dtype - sizeof(NVTEDType), // do_dtype - sizeof(NVTEDType), // dqkv_dtype + sizeof(NVTEDType), // qkv_dtype + sizeof(NVTEDType), // o_dtype + sizeof(NVTEDType), // do_dtype + sizeof(NVTEDType), // dqkv_dtype // data and scale layout - 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_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 // attention scaling - sizeof(float), // attn_scale + sizeof(float), // attn_scale // tensor dimensions - 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), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - sizeof(size_t), // num_tokens_q - sizeof(size_t), // num_tokens_kv + 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), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(size_t), // num_tokens_q + sizeof(size_t), // num_tokens_kv // paged KV dimensions - 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), // 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 // bias dimensions - 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), // bias_batch_size + sizeof(size_t), // bias_num_heads + sizeof(size_t), // bias_seqlen_q + sizeof(size_t), // bias_seqlen_kv }; bool operator<(const FusedAttnConfig &rhs) const { @@ -139,10 +139,10 @@ struct FusedAttnConfig { qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, attn_scale, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, - num_tokens_kv, bucketed_batch_size, bucketed_num_tokens_q, bucketed_num_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_batch_size, bias_num_heads, bias_seqlen_q, - bias_seqlen_kv) < + num_tokens_kv, bucketed_batch_size, bucketed_num_tokens_q, + bucketed_num_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_batch_size, bias_num_heads, + bias_seqlen_q, bias_seqlen_kv) < std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, @@ -216,39 +216,39 @@ struct FusedAttnFwdParams { cudaStream_t stream = nullptr; static constexpr size_t attr_sizes[] = { - sizeof(NVTETensor), // Q - sizeof(NVTETensor), // K - sizeof(NVTETensor), // V - sizeof(NVTETensor), // Bias - sizeof(NVTETensor), // SoftmaxOffset - sizeof(NVTETensor), // cu_seqlens_q - sizeof(NVTETensor), // cu_seqlens_kv - sizeof(NVTETensor), // cu_seqlens_q_padded - sizeof(NVTETensor), // cu_seqlens_kv_padded - sizeof(NVTETensor), // page_table_k - sizeof(NVTETensor), // page_table_v - sizeof(NVTETensor), // rng_state - sizeof(NVTETensor), // S - sizeof(NVTETensor), // O - sizeof(NVTETensorPack *), // Aux_CTX_Tensors - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - sizeof(NVTE_QKV_Layout), // qkv_layout - sizeof(NVTE_QKV_Format), // o_format - sizeof(NVTE_QKV_Format), // qkv_scale_inv_format - sizeof(NVTE_Bias_Type), // bias_type - sizeof(NVTE_Mask_Type), // attn_mask_type - sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(float), // attn_scale - sizeof(float), // dropout - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(uint8_t), // is_training - sizeof(uint8_t), // return_max_logit - sizeof(uint8_t), // cuda_graph - sizeof(NVTETensor), // workspace - sizeof(cudaStream_t), // stream + sizeof(NVTETensor), // Q + sizeof(NVTETensor), // K + sizeof(NVTETensor), // V + sizeof(NVTETensor), // Bias + sizeof(NVTETensor), // SoftmaxOffset + sizeof(NVTETensor), // cu_seqlens_q + sizeof(NVTETensor), // cu_seqlens_kv + sizeof(NVTETensor), // cu_seqlens_q_padded + sizeof(NVTETensor), // cu_seqlens_kv_padded + sizeof(NVTETensor), // page_table_k + sizeof(NVTETensor), // page_table_v + sizeof(NVTETensor), // rng_state + sizeof(NVTETensor), // S + sizeof(NVTETensor), // O + sizeof(NVTETensorPack *), // Aux_CTX_Tensors + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(float), // attn_scale + sizeof(float), // dropout + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(uint8_t), // is_training + sizeof(uint8_t), // return_max_logit + sizeof(uint8_t), // cuda_graph + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream }; }; diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 5737475bc9..baea567ad8 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -278,8 +278,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi 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_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); @@ -490,9 +490,9 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { input_cu_seqlens_kv_padded, input_page_table_k, input_page_table_v, input_rng_state, wkspace, p.stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, - p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, - wkspace, p.stream, handle); + fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, + output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, + input_rng_state, wkspace, p.stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } @@ -514,7 +514,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd); - transformer_engine::FusedAttnFwdParams p = transformer_engine::make_default_fused_attn_fwd_params(); + transformer_engine::FusedAttnFwdParams p = + transformer_engine::make_default_fused_attn_fwd_params(); p.Q = Q; p.K = K; p.V = V; @@ -639,12 +640,11 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); } - 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, p.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, p.stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { size_t i = 0; const Tensor *input_M = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); @@ -683,7 +683,8 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_bwd); - transformer_engine::FusedAttnBwdParams p = transformer_engine::make_default_fused_attn_bwd_params(); + transformer_engine::FusedAttnBwdParams p = + transformer_engine::make_default_fused_attn_bwd_params(); p.Q = Q; p.K = K; p.V = V; 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 c4918991d0..24fd65520d 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 @@ -1044,13 +1044,15 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } // namespace fused_attn using namespace transformer_engine::fused_attn; -void fused_attn_arbitrary_seqlen_fwd( - 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_fwd(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; @@ -1191,14 +1193,16 @@ void fused_attn_arbitrary_seqlen_fwd( } } -void fused_attn_arbitrary_seqlen_bwd( - 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) { +void fused_attn_arbitrary_seqlen_bwd(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 NVTE_Bias_Type bias_type = cfg.bias_type; 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 5065fbe93a..d570412e12 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 @@ -20,22 +20,26 @@ #include "transformer_engine/fused_attn.h" namespace transformer_engine { -void fused_attn_arbitrary_seqlen_fwd( - 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( - 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); +void fused_attn_arbitrary_seqlen_fwd(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(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; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 8da0ad2261..b2f4172767 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -16,13 +16,14 @@ namespace fused_attn { using namespace transformer_engine; // fused attention FWD FP8 with FE 1.0+ -void fused_attn_fp8_fwd_impl( - 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* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, - void* workspace, size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { +void fused_attn_fp8_fwd_impl(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* 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(); @@ -435,7 +436,7 @@ void fused_attn_fp8_fwd_impl( // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl( - const FusedAttnConfig &cfg, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, + 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, @@ -1070,11 +1071,12 @@ void fused_attn_fp8_bwd_impl( } // namespace fused_attn // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd( - 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) { +void fused_attn_fp8_fwd(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; @@ -1176,13 +1178,14 @@ void fused_attn_fp8_fwd( } } // fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd( - 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) { +void fused_attn_fp8_bwd(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; @@ -1268,11 +1271,12 @@ void fused_attn_fp8_bwd( fused_attn::fused_attn_fp8_bwd_impl( 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, 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"); } @@ -1290,7 +1294,7 @@ void fused_attn_fp8_bwd( } } -std::string is_supported_fp8_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { +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( @@ -1312,7 +1316,7 @@ std::string is_supported_fp8_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handl } } -std::string is_supported_fp8_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { +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( diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 1ede20c7f1..4193236215 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -16,20 +16,22 @@ namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd( - 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); +void fused_attn_fp8_fwd(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( - 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); +void fused_attn_fp8_bwd(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; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 36601d2b2c..df208a4337 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -335,8 +335,8 @@ void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); /*! \brief Query an attribute in a fused-attention forward-parameter object. */ void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, - NVTEFusedAttnFwdParamsAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written); + NVTEFusedAttnFwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); /*! \brief Set an attribute in a fused-attention forward-parameter object. */ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, @@ -1155,7 +1155,8 @@ class FusedAttnConfigWrapper { return *this; } FusedAttnConfigWrapper &set_bias_batch_size(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasBatchSize, &val, sizeof(val)); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasBatchSize, &val, + sizeof(val)); return *this; } FusedAttnConfigWrapper &set_bias_num_heads(size_t val) noexcept { @@ -1222,39 +1223,48 @@ class FusedAttnFwdParamsWrapper { return *this; } FusedAttnFwdParamsWrapper &set_Bias(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBias, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBias, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_SoftmaxOffset(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxOffset, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxOffset, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQ, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQ, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKV, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKV, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQPadded, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQPadded, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, + &val, sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_page_table_k(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableK, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableK, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_page_table_v(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableV, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableV, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_rng_state(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsRngState, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsRngState, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { @@ -1265,86 +1275,106 @@ class FusedAttnFwdParamsWrapper { nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsO, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack * val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAuxCtxTensors, &val, sizeof(val)); + FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack *val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAuxCtxTensors, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenQ, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenQ, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenKV, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenKV, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVLayout, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVLayout, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsOFormat, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsOFormat, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVScaleInvFormat, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + &val, sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBiasType, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBiasType, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnMaskType, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnMaskType, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxType, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxType, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnScale, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnScale, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsDropout, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsDropout, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeLeft, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeLeft, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeRight, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeRight, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBottomRightDiagonal, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBottomRightDiagonal, + &u8_val, sizeof(u8_val)); return *this; } FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsIsTraining, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsIsTraining, &u8_val, + sizeof(u8_val)); return *this; } FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsReturnMaxLogit, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsReturnMaxLogit, + &u8_val, sizeof(u8_val)); return *this; } FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCudaGraph, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCudaGraph, &u8_val, + sizeof(u8_val)); return *this; } FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWorkspace, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWorkspace, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsStream, &val, sizeof(val)); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsStream, &val, + sizeof(val)); return *this; } + private: NVTEFusedAttnFwdParams params_ = nullptr; }; @@ -1403,8 +1433,9 @@ class FusedAttnBwdParamsWrapper { nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDP, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(const NVTETensorPack * val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAuxCtxTensors, &val, sizeof(val)); + FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(const NVTETensorPack *val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAuxCtxTensors, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_dQ(NVTETensor val) noexcept { @@ -1420,112 +1451,139 @@ class FusedAttnBwdParamsWrapper { return *this; } FusedAttnBwdParamsWrapper &set_dBias(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDBias, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDBias, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_dSoftmaxOffset(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDSoftmaxOffset, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDSoftmaxOffset, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQ, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQ, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKV, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKV, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQPadded, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQPadded, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, + &val, sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenQ, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenQ, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenKV, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenKV, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVLayout, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVLayout, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsOFormat, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsOFormat, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOFormat, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOFormat, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQKVLayout, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQKVLayout, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, + &val, sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOScaleInvFormat, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOScaleInvFormat, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBiasType, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBiasType, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnMaskType, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnMaskType, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsSoftmaxType, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsSoftmaxType, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnScale, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnScale, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDropout, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDropout, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeLeft, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeLeft, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeRight, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeRight, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBottomRightDiagonal, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBottomRightDiagonal, + &u8_val, sizeof(u8_val)); return *this; } FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDeterministic, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDeterministic, &u8_val, + sizeof(u8_val)); return *this; } FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCudaGraph, &u8_val, sizeof(u8_val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCudaGraph, &u8_val, + sizeof(u8_val)); return *this; } FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWorkspace, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWorkspace, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsStream, &val, sizeof(val)); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsStream, &val, + sizeof(val)); return *this; } + private: NVTEFusedAttnBwdParams params_ = nullptr; }; diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 0d328b734d..6cc1fbb4d5 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -164,7 +164,9 @@ def get_fused_attn_backend(self): if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: bias_batch = self.bias_batch if self.bias_batch is not None else self.batch_size bias_heads = self.bias_heads if self.bias_heads is not None else self.q_num_heads - bias_seqlen_q = self.bias_seqlen_q if self.bias_seqlen_q is not None else self.q_max_seqlen + bias_seqlen_q = ( + self.bias_seqlen_q if self.bias_seqlen_q is not None else self.q_max_seqlen + ) bias_seqlen_kv = ( self.bias_seqlen_kv if self.bias_seqlen_kv is not None else self.kv_max_seqlen ) diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index d95b1db69f..30ff61b013 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -351,13 +351,12 @@ static void FusedAttnForwardImpl( 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, - bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); + 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, bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); 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) */ @@ -683,13 +682,12 @@ static void FusedAttnBackwardImpl( nvte_tensor_pack_create(&aux_input_tensors); 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, - bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); + 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, bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); 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/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 1551681dff..c293aeae88 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 @@ -1591,7 +1591,11 @@ def forward( attn_mask_type=attn_mask_type, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, - alibi_slopes_shape=alibi_slopes.shape if core_attention_bias_type == "alibi" and alibi_slopes is not None else None, + alibi_slopes_shape=( + alibi_slopes.shape + if core_attention_bias_type == "alibi" and alibi_slopes is not None + else None + ), core_attention_bias_type=core_attention_bias_type, core_attention_bias_shape=core_attention_bias_shape, core_attention_bias_requires_grad=( From c03d8523972423f9ba58aa6c51fc364a529a586c Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 15 Jul 2026 07:17:31 -0700 Subject: [PATCH 29/63] reorder struct fields, consolidate APIs to derive, make_config, make_cache_key, fix Jax bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.cpp | 484 +++++++++++------- .../common/fused_attn/config_and_params.h | 177 ++++--- .../common/fused_attn/fused_attn.cpp | 142 +---- .../fused_attn_f16_arbitrary_seqlen.cu | 12 +- .../fused_attn_f16_arbitrary_seqlen.h | 4 +- .../common/fused_attn/fused_attn_fp8.cu | 4 +- .../common/fused_attn/fused_attn_fp8.h | 4 +- transformer_engine/common/fused_attn/utils.cu | 2 - .../include/transformer_engine/fused_attn.h | 465 +++++++++-------- .../jax/cpp_extensions/attention.py | 52 +- .../attention/dot_product_attention/utils.py | 13 +- transformer_engine/pytorch/csrc/extensions.h | 4 +- .../pytorch/csrc/extensions/attention.cpp | 4 +- 13 files changed, 669 insertions(+), 698 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 1c7bb6ae4e..33088115af 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -10,6 +10,7 @@ #include +#include "../common.h" #include "../util/cuda_runtime.h" namespace { @@ -33,56 +34,52 @@ 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."); +void FusedAttnConfig::derive() { + const int64_t b = static_cast(batch_size); + const int64_t sq = static_cast(max_seqlen_q); + const int64_t skv = static_cast(max_seqlen_kv); - const int64_t b = static_cast(cfg->batch_size); - 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 NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); const bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); - 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); + const size_t tokens_q = num_tokens_q != 0 ? num_tokens_q : static_cast(b * sq); + const size_t tokens_kv = num_tokens_kv != 0 ? 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; + bucketed_batch_size = + (is_ragged_q || is_ragged_kv) ? fused_attn::get_max_batch_size(batch_size) : 0; + bucketed_num_tokens_q = is_ragged_q ? fused_attn::get_max_tokens(tokens_q) : 0; + bucketed_num_tokens_kv = is_ragged_kv ? fused_attn::get_max_tokens(tokens_kv) : 0; if (is_paged_kv) { - if (cfg->num_pages_k == 0) { - cfg->num_pages_k = static_cast(b); + if (num_pages_k == 0) { + num_pages_k = static_cast(b); } - if (cfg->num_pages_v == 0) { - cfg->num_pages_v = static_cast(b); + if (num_pages_v == 0) { + num_pages_v = static_cast(b); } - if (cfg->page_size_k == 0) { - cfg->page_size_k = static_cast(skv); + if (page_size_k == 0) { + page_size_k = static_cast(skv); } - if (cfg->page_size_v == 0) { - cfg->page_size_v = static_cast(skv); + if (page_size_v == 0) { + page_size_v = static_cast(skv); } - if (cfg->max_pages_per_seq_k == 0) { - cfg->max_pages_per_seq_k = 1; + if (max_pages_per_seq_k == 0) { + max_pages_per_seq_k = 1; } - if (cfg->max_pages_per_seq_v == 0) { - cfg->max_pages_per_seq_v = 1; + if (max_pages_per_seq_v == 0) { + max_pages_per_seq_v = 1; } } } -FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg, bool is_forward) { - FusedAttnConfig cache_cfg = cfg; +FusedAttnConfig FusedAttnConfig::make_cache_key(bool is_forward) const { + FusedAttnConfig cache_cfg = *this; const int64_t s_q = static_cast(cache_cfg.max_seqlen_q); const int64_t s_kv = static_cast(cache_cfg.max_seqlen_kv); @@ -139,8 +136,9 @@ FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg, b return cache_cfg; } -FusedAttnConfig make_fused_attn_config(const FusedAttnFwdParams ¶ms) { - FusedAttnConfig cfg = make_default_fused_attn_config(); +FusedAttnConfig FusedAttnFwdParams::make_config() const { + const FusedAttnFwdParams ¶ms = *this; + FusedAttnConfig cfg{}; cfg.is_training = params.is_training; cfg.deterministic = false; cfg.cuda_graph = params.cuda_graph; @@ -158,11 +156,93 @@ FusedAttnConfig make_fused_attn_config(const FusedAttnFwdParams ¶ms) { cfg.window_size_left = params.window_size_left; cfg.window_size_right = params.window_size_right; cfg.bottom_right_diagonal = params.bottom_right_diagonal; + + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(params.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(params.cu_seqlens_kv); + const Tensor *input_page_table_k = convertNVTETensorCheck(params.page_table_k); + const Tensor *input_page_table_v = convertNVTETensorCheck(params.page_table_v); + const Tensor *input_Q = convertNVTETensorCheck(params.Q); + const Tensor *input_K = convertNVTETensorCheck(params.K); + const Tensor *input_V = convertNVTETensorCheck(params.V); + const Tensor *input_Bias = convertNVTETensorCheck(params.Bias); + const Tensor *output_O = convertNVTETensorCheck(params.O); + + const NVTE_QKV_Format q_format = nvte_get_q_format(params.qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(params.qkv_layout); + auto *q_dims = input_Q->data.shape.data(); + auto *k_dims = input_K->data.shape.data(); + auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING + ? input_V->data.shape.data() + : input_V->columnwise_data.shape.data(); + AttentionShape q_shape(q_format, q_dims); + AttentionShape k_shape(kv_format, k_dims); + AttentionShape v_shape(kv_format, v_dims); + size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); + size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); + if (q_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_q->data.shape[0] - 1; + } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_kv->data.shape[0] - 1; + } + + int64_t num_pages_k = 0, num_pages_v = 0, page_size_k = 0, page_size_v = 0; + int64_t max_pages_per_seq_k = 0, max_pages_per_seq_v = 0; + if (input_page_table_k->data.dptr != nullptr) { + max_pages_per_seq_k = input_page_table_k->data.shape[1]; + } + if (input_page_table_v->data.dptr != nullptr) { + max_pages_per_seq_v = input_page_table_v->data.shape[1]; + } + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(params.qkv_layout); + if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { + const NVTE_QKV_Format paged_kv_format = nvte_get_kv_format(params.qkv_layout); + if (paged_kv_format == NVTE_QKV_Format::NVTE_BSHD) { + num_pages_k = input_K->data.shape[0]; + page_size_k = input_K->data.shape[1]; + num_pages_v = input_V->data.shape[0]; + page_size_v = input_V->data.shape[1]; + } else if (paged_kv_format == NVTE_QKV_Format::NVTE_SBHD) { + num_pages_k = input_K->data.shape[1]; + page_size_k = input_K->data.shape[0]; + num_pages_v = input_V->data.shape[1]; + page_size_v = input_V->data.shape[0]; + } + } + + 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."); + + cfg.scaling_mode = input_Q->scaling_mode; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = static_cast(output_O->data.dtype); + 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.num_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; + + if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && + input_Bias->data.dptr != nullptr && input_Bias->data.shape.size() >= 4) { + cfg.bias_batch_size = input_Bias->data.shape[0]; + cfg.bias_num_heads = input_Bias->data.shape[1]; + cfg.bias_seqlen_q = input_Bias->data.shape[2]; + cfg.bias_seqlen_kv = input_Bias->data.shape[3]; + } return cfg; } -FusedAttnConfig make_fused_attn_config(const FusedAttnBwdParams ¶ms) { - FusedAttnConfig cfg = make_default_fused_attn_config(); +FusedAttnConfig FusedAttnBwdParams::make_config() const { + const FusedAttnBwdParams ¶ms = *this; + FusedAttnConfig cfg{}; cfg.is_training = true; cfg.deterministic = params.deterministic; cfg.cuda_graph = params.cuda_graph; @@ -183,14 +263,64 @@ FusedAttnConfig make_fused_attn_config(const FusedAttnBwdParams ¶ms) { cfg.window_size_left = params.window_size_left; cfg.window_size_right = params.window_size_right; cfg.bottom_right_diagonal = params.bottom_right_diagonal; + + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(params.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(params.cu_seqlens_kv); + const Tensor *input_Q = convertNVTETensorCheck(params.Q); + const Tensor *input_K = convertNVTETensorCheck(params.K); + const Tensor *input_V = convertNVTETensorCheck(params.V); + const Tensor *input_O = convertNVTETensorCheck(params.O); + const Tensor *input_dO = convertNVTETensorCheck(params.dO); + const Tensor *output_dQ = convertNVTETensorCheck(params.dQ); + const Tensor *output_dBias = convertNVTETensorCheck(params.dBias); + + const NVTE_QKV_Format q_format = nvte_get_q_format(params.qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(params.qkv_layout); + auto *q_dims = input_Q->data.shape.data(); + auto *k_dims = input_K->data.shape.data(); + auto *v_dims = input_V->data.shape.data(); + AttentionShape q_shape(q_format, q_dims); + AttentionShape k_shape(kv_format, k_dims); + AttentionShape v_shape(kv_format, v_dims); + size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); + size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); + if (q_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_q->data.shape[0] - 1; + } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_kv->data.shape[0] - 1; + } + + 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."); + + cfg.scaling_mode = input_Q->scaling_mode; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = static_cast(input_O->data.dtype); + cfg.do_dtype = static_cast(input_dO->data.dtype); + cfg.dqkv_dtype = static_cast(output_dQ->data.dtype); + 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_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; + + if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && + output_dBias->data.shape.size() >= 4) { + cfg.bias_batch_size = output_dBias->data.shape[0]; + cfg.bias_num_heads = output_dBias->data.shape[1]; + cfg.bias_seqlen_q = output_dBias->data.shape[2]; + cfg.bias_seqlen_kv = output_dBias->data.shape[3]; + } return cfg; } } // namespace transformer_engine NVTEFusedAttnConfig nvte_create_fused_attn_config() { - return new transformer_engine::FusedAttnConfig( - transformer_engine::make_default_fused_attn_config()); + return new transformer_engine::FusedAttnConfig{}; } void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config) { @@ -254,6 +384,9 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigDropout: std::memcpy(buf, &cfg.dropout, attr_size); break; + case kNVTEFusedAttnConfigAttnScale: + std::memcpy(buf, &cfg.attn_scale, attr_size); + break; case kNVTEFusedAttnConfigQKVDtype: std::memcpy(buf, &cfg.qkv_dtype, attr_size); break; @@ -284,16 +417,13 @@ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigDOScaleInvFormat: std::memcpy(buf, &cfg.do_scale_inv_format, attr_size); break; - case kNVTEFusedAttnConfigAttnScale: - std::memcpy(buf, &cfg.attn_scale, 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: + case kNVTEFusedAttnConfigNumGQAGroups: std::memcpy(buf, &cfg.num_gqa_groups, attr_size); break; case kNVTEFusedAttnConfigHeadDimQK: @@ -401,6 +531,9 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigDropout: std::memcpy(&cfg.dropout, buf, attr_size); break; + case kNVTEFusedAttnConfigAttnScale: + std::memcpy(&cfg.attn_scale, buf, attr_size); + break; case kNVTEFusedAttnConfigQKVDtype: std::memcpy(&cfg.qkv_dtype, buf, attr_size); break; @@ -431,16 +564,13 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, case kNVTEFusedAttnConfigDOScaleInvFormat: std::memcpy(&cfg.do_scale_inv_format, buf, attr_size); break; - case kNVTEFusedAttnConfigAttnScale: - std::memcpy(&cfg.attn_scale, 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: + case kNVTEFusedAttnConfigNumGQAGroups: std::memcpy(&cfg.num_gqa_groups, buf, attr_size); break; case kNVTEFusedAttnConfigHeadDimQK: @@ -497,24 +627,13 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, } NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params() { - return new transformer_engine::FusedAttnFwdParams( - transformer_engine::make_default_fused_attn_fwd_params()); + return new transformer_engine::FusedAttnFwdParams{}; } void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { delete transformer_engine::get_fused_attn_fwd_params_mutable(params); } -#define NVTE_FWD_PARAMS_GET_BOOL_FIELD(ATTR, FIELD) \ - case ATTR: \ - bool_to_uint8(p.FIELD, buf); \ - break - -#define NVTE_FWD_PARAMS_SET_BOOL_FIELD(ATTR, FIELD) \ - case ATTR: \ - uint8_to_bool(buf, p.FIELD); \ - break - void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, NVTEFusedAttnFwdParamsAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written) { @@ -577,47 +696,54 @@ void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, case kNVTEFusedAttnFwdParamsAuxCtxTensors: std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); break; - case kNVTEFusedAttnFwdParamsMaxSeqlenQ: - std::memcpy(buf, &p.max_seqlen_q, attr_size); - break; - case kNVTEFusedAttnFwdParamsMaxSeqlenKV: - std::memcpy(buf, &p.max_seqlen_kv, attr_size); + case kNVTEFusedAttnFwdParamsIsTraining: + bool_to_uint8(p.is_training, buf); break; - case kNVTEFusedAttnFwdParamsQKVLayout: - std::memcpy(buf, &p.qkv_layout, attr_size); + case kNVTEFusedAttnFwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); break; - case kNVTEFusedAttnFwdParamsOFormat: - std::memcpy(buf, &p.o_format, attr_size); + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + bool_to_uint8(p.return_max_logit, buf); break; - case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: - std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); break; case kNVTEFusedAttnFwdParamsBiasType: std::memcpy(buf, &p.bias_type, attr_size); break; - case kNVTEFusedAttnFwdParamsAttnMaskType: - std::memcpy(buf, &p.attn_mask_type, attr_size); - break; case kNVTEFusedAttnFwdParamsSoftmaxType: std::memcpy(buf, &p.softmax_type, attr_size); break; - case kNVTEFusedAttnFwdParamsAttnScale: - std::memcpy(buf, &p.attn_scale, attr_size); - break; - case kNVTEFusedAttnFwdParamsDropout: - std::memcpy(buf, &p.dropout, attr_size); - break; case kNVTEFusedAttnFwdParamsWindowSizeLeft: std::memcpy(buf, &p.window_size_left, attr_size); break; case kNVTEFusedAttnFwdParamsWindowSizeRight: std::memcpy(buf, &p.window_size_right, attr_size); break; - NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsBottomRightDiagonal, - bottom_right_diagonal); - NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsIsTraining, is_training); - NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsReturnMaxLogit, return_max_logit); - NVTE_FWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnFwdParamsCudaGraph, cuda_graph); + case kNVTEFusedAttnFwdParamsBottomRightDiagonal: + bool_to_uint8(p.bottom_right_diagonal, buf); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; case kNVTEFusedAttnFwdParamsWorkspace: std::memcpy(buf, &p.workspace, attr_size); break; @@ -686,47 +812,54 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, case kNVTEFusedAttnFwdParamsAuxCtxTensors: std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); break; - case kNVTEFusedAttnFwdParamsMaxSeqlenQ: - std::memcpy(&p.max_seqlen_q, buf, attr_size); + case kNVTEFusedAttnFwdParamsIsTraining: + uint8_to_bool(buf, p.is_training); break; - case kNVTEFusedAttnFwdParamsMaxSeqlenKV: - std::memcpy(&p.max_seqlen_kv, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsQKVLayout: - std::memcpy(&p.qkv_layout, buf, attr_size); + case kNVTEFusedAttnFwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); break; - case kNVTEFusedAttnFwdParamsOFormat: - std::memcpy(&p.o_format, buf, attr_size); + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + uint8_to_bool(buf, p.return_max_logit); break; - case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: - std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); break; case kNVTEFusedAttnFwdParamsBiasType: std::memcpy(&p.bias_type, buf, attr_size); break; - case kNVTEFusedAttnFwdParamsAttnMaskType: - std::memcpy(&p.attn_mask_type, buf, attr_size); - break; case kNVTEFusedAttnFwdParamsSoftmaxType: std::memcpy(&p.softmax_type, buf, attr_size); break; - case kNVTEFusedAttnFwdParamsAttnScale: - std::memcpy(&p.attn_scale, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsDropout: - std::memcpy(&p.dropout, buf, attr_size); - break; case kNVTEFusedAttnFwdParamsWindowSizeLeft: std::memcpy(&p.window_size_left, buf, attr_size); break; case kNVTEFusedAttnFwdParamsWindowSizeRight: std::memcpy(&p.window_size_right, buf, attr_size); break; - NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsBottomRightDiagonal, - bottom_right_diagonal); - NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsIsTraining, is_training); - NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsReturnMaxLogit, return_max_logit); - NVTE_FWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnFwdParamsCudaGraph, cuda_graph); + case kNVTEFusedAttnFwdParamsBottomRightDiagonal: + uint8_to_bool(buf, p.bottom_right_diagonal); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; case kNVTEFusedAttnFwdParamsWorkspace: std::memcpy(&p.workspace, buf, attr_size); break; @@ -738,28 +871,14 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, } } -#undef NVTE_FWD_PARAMS_GET_BOOL_FIELD -#undef NVTE_FWD_PARAMS_SET_BOOL_FIELD - NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params() { - return new transformer_engine::FusedAttnBwdParams( - transformer_engine::make_default_fused_attn_bwd_params()); + return new transformer_engine::FusedAttnBwdParams{}; } void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { delete transformer_engine::get_fused_attn_bwd_params_mutable(params); } -#define NVTE_BWD_PARAMS_GET_BOOL_FIELD(ATTR, FIELD) \ - case ATTR: \ - bool_to_uint8(p.FIELD, buf); \ - break - -#define NVTE_BWD_PARAMS_SET_BOOL_FIELD(ATTR, FIELD) \ - case ATTR: \ - uint8_to_bool(buf, p.FIELD); \ - break - void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written) { @@ -828,11 +947,35 @@ void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenQ: - std::memcpy(buf, &p.max_seqlen_q, attr_size); + case kNVTEFusedAttnBwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenKV: - std::memcpy(buf, &p.max_seqlen_kv, attr_size); + case kNVTEFusedAttnBwdParamsDeterministic: + bool_to_uint8(p.deterministic, buf); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeLeft: + std::memcpy(buf, &p.window_size_left, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeRight: + std::memcpy(buf, &p.window_size_right, attr_size); + break; + case kNVTEFusedAttnBwdParamsBottomRightDiagonal: + bool_to_uint8(p.bottom_right_diagonal, buf); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); break; case kNVTEFusedAttnBwdParamsQKVLayout: std::memcpy(buf, &p.qkv_layout, attr_size); @@ -852,31 +995,12 @@ void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsDOScaleInvFormat: std::memcpy(buf, &p.do_scale_inv_format, attr_size); break; - case kNVTEFusedAttnBwdParamsBiasType: - std::memcpy(buf, &p.bias_type, attr_size); - break; - case kNVTEFusedAttnBwdParamsAttnMaskType: - std::memcpy(buf, &p.attn_mask_type, attr_size); - break; - case kNVTEFusedAttnBwdParamsSoftmaxType: - std::memcpy(buf, &p.softmax_type, attr_size); - break; - case kNVTEFusedAttnBwdParamsAttnScale: - std::memcpy(buf, &p.attn_scale, attr_size); - break; - case kNVTEFusedAttnBwdParamsDropout: - std::memcpy(buf, &p.dropout, attr_size); - break; - case kNVTEFusedAttnBwdParamsWindowSizeLeft: - std::memcpy(buf, &p.window_size_left, attr_size); + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); break; - case kNVTEFusedAttnBwdParamsWindowSizeRight: - std::memcpy(buf, &p.window_size_right, attr_size); + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); break; - NVTE_BWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnBwdParamsBottomRightDiagonal, - bottom_right_diagonal); - NVTE_BWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnBwdParamsDeterministic, deterministic); - NVTE_BWD_PARAMS_GET_BOOL_FIELD(kNVTEFusedAttnBwdParamsCudaGraph, cuda_graph); case kNVTEFusedAttnBwdParamsWorkspace: std::memcpy(buf, &p.workspace, attr_size); break; @@ -951,11 +1075,35 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenQ: - std::memcpy(&p.max_seqlen_q, buf, attr_size); + case kNVTEFusedAttnBwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenKV: - std::memcpy(&p.max_seqlen_kv, buf, attr_size); + case kNVTEFusedAttnBwdParamsDeterministic: + uint8_to_bool(buf, p.deterministic); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(&p.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeLeft: + std::memcpy(&p.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeRight: + std::memcpy(&p.window_size_right, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsBottomRightDiagonal: + uint8_to_bool(buf, p.bottom_right_diagonal); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); break; case kNVTEFusedAttnBwdParamsQKVLayout: std::memcpy(&p.qkv_layout, buf, attr_size); @@ -975,31 +1123,12 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsDOScaleInvFormat: std::memcpy(&p.do_scale_inv_format, buf, attr_size); break; - case kNVTEFusedAttnBwdParamsBiasType: - std::memcpy(&p.bias_type, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsAttnMaskType: - std::memcpy(&p.attn_mask_type, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsSoftmaxType: - std::memcpy(&p.softmax_type, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsAttnScale: - std::memcpy(&p.attn_scale, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDropout: - std::memcpy(&p.dropout, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsWindowSizeLeft: - std::memcpy(&p.window_size_left, buf, attr_size); + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); break; - case kNVTEFusedAttnBwdParamsWindowSizeRight: - std::memcpy(&p.window_size_right, buf, attr_size); + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); break; - NVTE_BWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnBwdParamsBottomRightDiagonal, - bottom_right_diagonal); - NVTE_BWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnBwdParamsDeterministic, deterministic); - NVTE_BWD_PARAMS_SET_BOOL_FIELD(kNVTEFusedAttnBwdParamsCudaGraph, cuda_graph); case kNVTEFusedAttnBwdParamsWorkspace: std::memcpy(&p.workspace, buf, attr_size); break; @@ -1010,6 +1139,3 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); } } - -#undef NVTE_BWD_PARAMS_GET_BOOL_FIELD -#undef NVTE_BWD_PARAMS_SET_BOOL_FIELD diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 5f31ecb0ed..1305f21b2f 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -20,7 +20,7 @@ namespace transformer_engine { struct FusedAttnConfig { // basic attention knobs - bool is_training = false; + bool is_training = true; bool deterministic = false; bool cuda_graph = false; bool return_max_logit = false; @@ -32,6 +32,7 @@ struct FusedAttnConfig { NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; float dropout = 0.0f; + float attn_scale = 1.0f; // data types NVTEDType qkv_dtype = kNVTEBFloat16; @@ -47,9 +48,6 @@ struct FusedAttnConfig { NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; - // attention scaling - float attn_scale = 1.0f; - // tensor dimensions size_t batch_size = 0; size_t num_attn_heads = 0; @@ -94,6 +92,7 @@ struct FusedAttnConfig { sizeof(NVTE_Softmax_Type), // softmax_type sizeof(NVTEScalingMode), // scaling_mode sizeof(float), // dropout + sizeof(float), // attn_scale // data types sizeof(NVTEDType), // qkv_dtype sizeof(NVTEDType), // o_dtype @@ -106,8 +105,6 @@ struct FusedAttnConfig { sizeof(NVTE_QKV_Layout), // dqkv_layout sizeof(NVTE_QKV_Format), // qkv_scale_inv_format sizeof(NVTE_QKV_Format), // do_scale_inv_format - // attention scaling - sizeof(float), // attn_scale // tensor dimensions sizeof(size_t), // batch_size sizeof(size_t), // num_attn_heads @@ -135,9 +132,9 @@ struct FusedAttnConfig { bool operator<(const FusedAttnConfig &rhs) const { return std::tie(is_training, deterministic, cuda_graph, return_max_logit, attn_mask_type, bias_type, window_size_left, window_size_right, bottom_right_diagonal, - softmax_type, scaling_mode, dropout, qkv_dtype, o_dtype, do_dtype, dqkv_dtype, - qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, attn_scale, batch_size, num_attn_heads, num_gqa_groups, + softmax_type, scaling_mode, dropout, attn_scale, qkv_dtype, o_dtype, do_dtype, + dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_tokens_kv, bucketed_batch_size, bucketed_num_tokens_q, bucketed_num_tokens_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, @@ -146,9 +143,10 @@ struct FusedAttnConfig { std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, - rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, rhs.qkv_layout, - rhs.o_format, rhs.do_format, rhs.dqkv_layout, rhs.qkv_scale_inv_format, - rhs.do_scale_inv_format, rhs.attn_scale, rhs.batch_size, rhs.num_attn_heads, + rhs.attn_scale, rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, + rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, + rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, + rhs.num_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, rhs.bucketed_batch_size, rhs.bucketed_num_tokens_q, rhs.bucketed_num_tokens_kv, rhs.num_pages_k, @@ -156,19 +154,17 @@ struct FusedAttnConfig { rhs.max_pages_per_seq_v, rhs.bias_batch_size, rhs.bias_num_heads, rhs.bias_seqlen_q, rhs.bias_seqlen_kv); } -}; - -inline FusedAttnConfig make_default_fused_attn_config() { return FusedAttnConfig{}; } -void populate_fused_attn_config(FusedAttnConfig *cfg); + // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields + // that have been set by the caller. + void derive(); -// 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(). Pass is_forward=true when -// keying a forward graph and is_forward=false for a backward graph; each key drops the fields the -// corresponding graph does not consume so it is not fragmented by them: a training forward key -// drops the dO/dQKV dtypes and the (backward-only) deterministic flag, and a backward key drops -// the (forward-only) return_max_logit flag. -FusedAttnConfig make_fused_attn_graph_cache_config(const FusedAttnConfig &cfg, bool is_forward); + // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. + // It drops fields that are invariant (e.g. batch_size) or irrelevant (e.g. dO/dQKV dtypes + // and `deterministic` for forward, and `return_max_logit` for backward) to the corresponding graph. + // This helps avoid redundant graph builds and cache misses. + FusedAttnConfig make_cache_key(bool is_forward) const; +}; inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); @@ -196,22 +192,22 @@ struct FusedAttnFwdParams { NVTETensor S = nullptr; NVTETensor O = nullptr; NVTETensorPack *Aux_CTX_Tensors = nullptr; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; - NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; - NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; - NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + bool is_training = true; + bool cuda_graph = false; + bool return_max_logit = false; NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; - float attn_scale = 1.0f; - float dropout = 0.0f; int64_t window_size_left = -1; int64_t window_size_right = -1; bool bottom_right_diagonal = true; - bool is_training = false; - bool return_max_logit = false; - bool cuda_graph = false; + float dropout = 0.0f; + float attn_scale = 1.0f; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; NVTETensor workspace = nullptr; cudaStream_t stream = nullptr; @@ -231,27 +227,43 @@ struct FusedAttnFwdParams { sizeof(NVTETensor), // S sizeof(NVTETensor), // O sizeof(NVTETensorPack *), // Aux_CTX_Tensors - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - sizeof(NVTE_QKV_Layout), // qkv_layout - sizeof(NVTE_QKV_Format), // o_format - sizeof(NVTE_QKV_Format), // qkv_scale_inv_format - sizeof(NVTE_Bias_Type), // bias_type + sizeof(uint8_t), // is_training + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(float), // attn_scale - sizeof(float), // dropout sizeof(int64_t), // window_size_left sizeof(int64_t), // window_size_right sizeof(uint8_t), // bottom_right_diagonal - sizeof(uint8_t), // is_training - sizeof(uint8_t), // return_max_logit - sizeof(uint8_t), // cuda_graph + sizeof(float), // dropout + sizeof(float), // attn_scale + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv sizeof(NVTETensor), // workspace sizeof(cudaStream_t), // stream }; + + // Build a FusedAttnConfig from the scalar "knobs" carried here (e.g. attn_mask_type, bias_type) + // and the fields derived from the tensor handles (dtypes, dims, scaling mode, paged-KV and bias + // broadcast shapes). Returns the real execution config; call FusedAttnConfig::make_cache_key on + // it to obtain the normalized cuDNN graph-cache key. + FusedAttnConfig make_config() const; }; +inline const FusedAttnFwdParams *get_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); + return reinterpret_cast(params); +} + +inline FusedAttnFwdParams *get_fused_attn_fwd_params_mutable(NVTEFusedAttnFwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); + return reinterpret_cast(params); +} + struct FusedAttnBwdParams { NVTETensor Q = nullptr; NVTETensor K = nullptr; @@ -270,24 +282,24 @@ struct FusedAttnBwdParams { NVTETensor cu_seqlens_kv = nullptr; NVTETensor cu_seqlens_q_padded = nullptr; NVTETensor cu_seqlens_kv_padded = nullptr; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; + bool cuda_graph = false; + bool deterministic = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + float dropout = 0.0f; + float attn_scale = 1.0f; 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; - float attn_scale = 1.0f; - float dropout = 0.0f; - int64_t window_size_left = -1; - int64_t window_size_right = -1; - bool bottom_right_diagonal = true; - bool deterministic = false; - bool cuda_graph = false; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; NVTETensor workspace = nullptr; cudaStream_t stream = nullptr; @@ -309,49 +321,34 @@ struct FusedAttnBwdParams { sizeof(NVTETensor), // cu_seqlens_kv sizeof(NVTETensor), // cu_seqlens_q_padded sizeof(NVTETensor), // cu_seqlens_kv_padded - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // deterministic + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(float), // dropout + sizeof(float), // attn_scale 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(float), // attn_scale - sizeof(float), // dropout - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(uint8_t), // deterministic - sizeof(uint8_t), // cuda_graph + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv sizeof(NVTETensor), // workspace sizeof(cudaStream_t), // stream }; -}; - -inline FusedAttnFwdParams make_default_fused_attn_fwd_params() { return FusedAttnFwdParams{}; } - -inline FusedAttnBwdParams make_default_fused_attn_bwd_params() { return FusedAttnBwdParams{}; } -// Build a FusedAttnConfig from the scalar "knobs" carried by the fwd/bwd params (mask/bias/softmax -// type, scales, dropout, window, layout/format fields, flags). The tensor-derived fields (dtypes, -// dims, scaling_mode, paged-KV / bias dims, token counts) are left at their defaults and must be -// filled in by the caller from the actual tensors. -FusedAttnConfig make_fused_attn_config(const FusedAttnFwdParams ¶ms); -FusedAttnConfig make_fused_attn_config(const FusedAttnBwdParams ¶ms); - -inline const FusedAttnFwdParams *get_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { - NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); - return reinterpret_cast(params); -} - -inline FusedAttnFwdParams *get_fused_attn_fwd_params_mutable(NVTEFusedAttnFwdParams params) { - NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); - return reinterpret_cast(params); -} + // Build a FusedAttnConfig from the scalar "knobs" carried here (e.g. attn_mask_type, bias_type) + // and the fields derived from the tensor handles (e.g. dtypes, dims, scaling mode and bias broadcast + // shape). Returns the real execution config; call FusedAttnConfig::make_cache_key on it to + // obtain the normalized cuDNN graph-cache key. + FusedAttnConfig make_config() const; +}; inline const FusedAttnBwdParams *get_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index baea567ad8..c74088713b 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -348,7 +348,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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) { - transformer_engine::FusedAttnConfig cfg = transformer_engine::make_default_fused_attn_config(); + transformer_engine::FusedAttnConfig cfg{}; cfg.qkv_layout = qkv_layout; cfg.bias_type = bias_type; cfg.attn_mask_type = attn_mask_type; @@ -396,89 +396,8 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { Tensor *output_O = convertNVTETensorCheck(p.O); Tensor *wkspace = convertNVTETensor(p.workspace); - NVTE_QKV_Format q_format = nvte_get_q_format(p.qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(p.qkv_layout); - auto *q_dims = input_Q->data.shape.data(); - auto *k_dims = input_K->data.shape.data(); - auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING - ? input_V->data.shape.data() - : input_V->columnwise_data.shape.data(); - AttentionShape q_shape(q_format, q_dims); - AttentionShape k_shape(kv_format, k_dims); - AttentionShape v_shape(kv_format, v_dims); - size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); - size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); - if (q_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_q->data.shape[0] - 1; - } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_kv->data.shape[0] - 1; - } - - int64_t num_pages_k = 0; - int64_t num_pages_v = 0; - int64_t page_size_k = 0; - int64_t page_size_v = 0; - int64_t max_pages_per_seq_k = 0; - int64_t max_pages_per_seq_v = 0; - if (input_page_table_k->data.dptr != nullptr) { - max_pages_per_seq_k = input_page_table_k->data.shape[1]; - } - if (input_page_table_v->data.dptr != nullptr) { - max_pages_per_seq_v = input_page_table_v->data.shape[1]; - } - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(p.qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { - NVTE_QKV_Format paged_kv_format = nvte_get_kv_format(p.qkv_layout); - if (paged_kv_format == NVTE_QKV_Format::NVTE_BSHD) { - num_pages_k = input_K->data.shape[0]; - page_size_k = input_K->data.shape[1]; - num_pages_v = input_V->data.shape[0]; - page_size_v = input_V->data.shape[1]; - } else if (paged_kv_format == NVTE_QKV_Format::NVTE_SBHD) { - num_pages_k = input_K->data.shape[1]; - page_size_k = input_K->data.shape[0]; - num_pages_v = input_V->data.shape[1]; - page_size_v = input_V->data.shape[0]; - } - } - 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 ((p.bias_type != NVTE_NO_BIAS) && (p.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]; - } - - FusedAttnConfig cfg = make_fused_attn_config(p); - cfg.scaling_mode = scaling_mode; - 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; + FusedAttnConfig cfg = p.make_config(); NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), /*message=*/nullptr); @@ -514,8 +433,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd); - transformer_engine::FusedAttnFwdParams p = - transformer_engine::make_default_fused_attn_fwd_params(); + transformer_engine::FusedAttnFwdParams p{}; p.Q = Q; p.K = K; p.V = V; @@ -574,57 +492,8 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(p.dSoftmaxOffset); Tensor *wkspace = convertNVTETensor(p.workspace); - NVTE_QKV_Format q_format = nvte_get_q_format(p.qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(p.qkv_layout); - auto *q_dims = input_Q->data.shape.data(); - auto *k_dims = input_K->data.shape.data(); - auto *v_dims = input_V->data.shape.data(); - AttentionShape q_shape(q_format, q_dims); - AttentionShape k_shape(kv_format, k_dims); - AttentionShape v_shape(kv_format, v_dims); - size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); - size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); - if (q_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_q->data.shape[0] - 1; - } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_kv->data.shape[0] - 1; - } - 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 ((p.bias_type != NVTE_NO_BIAS) && (p.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]; - } - - FusedAttnConfig cfg = make_fused_attn_config(p); - cfg.scaling_mode = scaling_mode; - 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; + FusedAttnConfig cfg = p.make_config(); NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), /*message=*/nullptr); @@ -683,8 +552,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_bwd); - transformer_engine::FusedAttnBwdParams p = - transformer_engine::make_default_fused_attn_bwd_params(); + transformer_engine::FusedAttnBwdParams p{}; p.Q = Q; p.K = K; p.V = V; 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 24fd65520d..1fa7c3e9b8 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 @@ -142,7 +142,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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, /*is_forward=*/true); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(/*is_forward=*/true); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -625,7 +625,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( // 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, /*is_forward=*/false); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(/*is_forward=*/false); try { namespace fe = cudnn_frontend; @@ -1096,7 +1096,7 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i void *devPtrPageTableV = page_table_v ? page_table_v->data.dptr : nullptr; FusedAttnConfig graph_cfg = cfg; - populate_fused_attn_config(&graph_cfg); + graph_cfg.derive(); size_t i = 0; if (Aux_CTX_Tensors->size == 0) { @@ -1221,7 +1221,7 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i } FusedAttnConfig graph_cfg = cfg; - populate_fused_attn_config(&graph_cfg); + graph_cfg.derive(); void *devPtrdQ = output_dQ->data.dptr; void *devPtrdK = output_dK->data.dptr; @@ -1270,7 +1270,7 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; - populate_fused_attn_config(&graph_cfg); + graph_cfg.derive(); size_t workspace_size = 0; try { @@ -1294,7 +1294,7 @@ std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handl std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; - populate_fused_attn_config(&graph_cfg); + graph_cfg.derive(); size_t workspace_size = 0; try { 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 d570412e12..3f22f131d8 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 @@ -43,12 +43,12 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i // 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. +// if not, return a diagnostic message explaining why it is not supported. 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. +// if not, return a diagnostic message explaining why it is not supported. std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index b2f4172767..95f2eaebb4 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -82,7 +82,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de 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, /*is_forward=*/true); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(/*is_forward=*/true); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -515,7 +515,7 @@ 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, /*is_forward=*/false); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(/*is_forward=*/false); try { namespace fe = cudnn_frontend; using graph_and_tensors = diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 4193236215..f75906cbad 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -35,11 +35,11 @@ void fused_attn_fp8_bwd(const FusedAttnConfig &cfg, const Tensor *input_Q, const // 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. +// if not, return a diagnostic message explaining why it is not supported. 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. +// if not, return a diagnostic message explaining why it is not supported. 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 44413b40ef..34344335a6 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -10,7 +10,6 @@ #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" @@ -635,7 +634,6 @@ __global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, i } } // namespace fused_attn - } // 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/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index df208a4337..3010e6418a 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -200,14 +200,14 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); typedef void *NVTEFusedAttnConfig; /*! \enum NVTEFusedAttnConfigAttribute - * \brief Attribute types for ``NVTEFusedAttnConfig``. + * \brief Attributes for ``NVTEFusedAttnConfig``. * * This enum is used to index the ``FusedAttnConfig`` struct. The order of its fields must match that of - * the declaration fields and the ``attr_sizes`` array of that struct. New fields may only be appended - * at the end, and existing fields are never to be reordered, removed, or resized. + * the declaration fields and ``attr_sizes`` array of ``FusedAttnConfig``. New fields may only be appended + * at the end and existing fields are never reordered, removed, or resized. */ enum NVTEFusedAttnConfigAttribute { - // basic attention knobs + // basic configuration knobs kNVTEFusedAttnConfigIsTraining = 0, kNVTEFusedAttnConfigDeterministic, kNVTEFusedAttnConfigCudaGraph, @@ -220,24 +220,23 @@ enum NVTEFusedAttnConfigAttribute { kNVTEFusedAttnConfigSoftmaxType, kNVTEFusedAttnConfigScalingMode, kNVTEFusedAttnConfigDropout, - // data types + kNVTEFusedAttnConfigAttnScale, + // tensor types kNVTEFusedAttnConfigQKVDtype, kNVTEFusedAttnConfigODtype, kNVTEFusedAttnConfigDODtype, kNVTEFusedAttnConfigDQKVDtype, - // data and scale layout + // tensor layouts kNVTEFusedAttnConfigQKVLayout, kNVTEFusedAttnConfigOFormat, kNVTEFusedAttnConfigDOFormat, kNVTEFusedAttnConfigDQKVLayout, kNVTEFusedAttnConfigQKVScaleInvFormat, kNVTEFusedAttnConfigDOScaleInvFormat, - // attention scaling - kNVTEFusedAttnConfigAttnScale, // tensor dimensions kNVTEFusedAttnConfigBatchSize, kNVTEFusedAttnConfigNumAttnHeads, - kNVTEFusedAttnConfigNumGqaGroups, + kNVTEFusedAttnConfigNumGQAGroups, kNVTEFusedAttnConfigHeadDimQK, kNVTEFusedAttnConfigHeadDimV, kNVTEFusedAttnConfigMaxSeqlenQ, @@ -256,22 +255,14 @@ enum NVTEFusedAttnConfigAttribute { kNVTEFusedAttnConfigBiasNumHeads, kNVTEFusedAttnConfigBiasSeqlenQ, kNVTEFusedAttnConfigBiasSeqlenKV, + // number of attributes 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()``. - */ +/*! \brief Create a fused-attention configuration. */ NVTEFusedAttnConfig nvte_create_fused_attn_config(void); -/*! \brief Destroy a fused-attention configuration handle. */ +/*! \brief Destroy a fused-attention configuration. */ void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config); /*! \brief Query an attribute in a fused-attention configuration. */ @@ -288,9 +279,14 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, typedef void *NVTEFusedAttnFwdParams; /*! \enum NVTEFusedAttnFwdParamsAttribute - * \brief Attribute types for ``NVTEFusedAttnFwdParams``. + * \brief Attributes for ``NVTEFusedAttnFwdParams``. + * + * This enum is used to index the ``FusedAttnFwdParams`` struct. The order of its fields must match that of + * the declaration fields and ``attr_sizes`` array of ``FusedAttnFwdParams``. New fields may only be appended + * at the end and existing fields are never reordered, removed, or resized. */ enum NVTEFusedAttnFwdParamsAttribute { + // tensor handles kNVTEFusedAttnFwdParamsQ = 0, kNVTEFusedAttnFwdParamsK, kNVTEFusedAttnFwdParamsV, @@ -306,31 +302,34 @@ enum NVTEFusedAttnFwdParamsAttribute { kNVTEFusedAttnFwdParamsS, kNVTEFusedAttnFwdParamsO, kNVTEFusedAttnFwdParamsAuxCtxTensors, - kNVTEFusedAttnFwdParamsMaxSeqlenQ, - kNVTEFusedAttnFwdParamsMaxSeqlenKV, - kNVTEFusedAttnFwdParamsQKVLayout, - kNVTEFusedAttnFwdParamsOFormat, - kNVTEFusedAttnFwdParamsQKVScaleInvFormat, - kNVTEFusedAttnFwdParamsBiasType, + // configuration knobs + kNVTEFusedAttnFwdParamsIsTraining, + kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsReturnMaxLogit, kNVTEFusedAttnFwdParamsAttnMaskType, + kNVTEFusedAttnFwdParamsBiasType, kNVTEFusedAttnFwdParamsSoftmaxType, - kNVTEFusedAttnFwdParamsAttnScale, - kNVTEFusedAttnFwdParamsDropout, kNVTEFusedAttnFwdParamsWindowSizeLeft, kNVTEFusedAttnFwdParamsWindowSizeRight, kNVTEFusedAttnFwdParamsBottomRightDiagonal, - kNVTEFusedAttnFwdParamsIsTraining, - kNVTEFusedAttnFwdParamsReturnMaxLogit, - kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsDropout, + kNVTEFusedAttnFwdParamsAttnScale, + kNVTEFusedAttnFwdParamsQKVLayout, + kNVTEFusedAttnFwdParamsOFormat, + kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + kNVTEFusedAttnFwdParamsMaxSeqlenQ, + kNVTEFusedAttnFwdParamsMaxSeqlenKV, + // workspace and stream kNVTEFusedAttnFwdParamsWorkspace, kNVTEFusedAttnFwdParamsStream, + // number of attributes kNVTEFusedAttnFwdParamsNumAttributes }; -/*! \brief Create a default-initialized fused-attention forward-parameter object. */ +/*! \brief Create a fused-attention forward-parameter object. */ NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void); -/*! \brief Destroy a fused-attention forward-parameter handle. */ +/*! \brief Destroy a fused-attention forward-parameter object. */ void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); /*! \brief Query an attribute in a fused-attention forward-parameter object. */ @@ -347,9 +346,14 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, typedef void *NVTEFusedAttnBwdParams; /*! \enum NVTEFusedAttnBwdParamsAttribute - * \brief Attribute types for ``NVTEFusedAttnBwdParams``. + * \brief Attributes for ``NVTEFusedAttnBwdParams``. + * + * This enum is used to index the ``FusedAttnBwdParams`` struct. The order of its fields must match that of + * the declaration fields and ``attr_sizes`` array of ``FusedAttnBwdParams``. New fields may only be appended + * at the end and existing fields are never reordered, removed, or resized. */ enum NVTEFusedAttnBwdParamsAttribute { + // tensor handles kNVTEFusedAttnBwdParamsQ = 0, kNVTEFusedAttnBwdParamsK, kNVTEFusedAttnBwdParamsV, @@ -367,33 +371,36 @@ enum NVTEFusedAttnBwdParamsAttribute { kNVTEFusedAttnBwdParamsCuSeqlensKV, kNVTEFusedAttnBwdParamsCuSeqlensQPadded, kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, - kNVTEFusedAttnBwdParamsMaxSeqlenQ, - kNVTEFusedAttnBwdParamsMaxSeqlenKV, + // configuration knobs + kNVTEFusedAttnBwdParamsCudaGraph, + kNVTEFusedAttnBwdParamsDeterministic, + kNVTEFusedAttnBwdParamsAttnMaskType, + kNVTEFusedAttnBwdParamsBiasType, + kNVTEFusedAttnBwdParamsSoftmaxType, + kNVTEFusedAttnBwdParamsWindowSizeLeft, + kNVTEFusedAttnBwdParamsWindowSizeRight, + kNVTEFusedAttnBwdParamsBottomRightDiagonal, + kNVTEFusedAttnBwdParamsDropout, + kNVTEFusedAttnBwdParamsAttnScale, kNVTEFusedAttnBwdParamsQKVLayout, kNVTEFusedAttnBwdParamsOFormat, kNVTEFusedAttnBwdParamsDOFormat, kNVTEFusedAttnBwdParamsDQKVLayout, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, kNVTEFusedAttnBwdParamsDOScaleInvFormat, - kNVTEFusedAttnBwdParamsBiasType, - kNVTEFusedAttnBwdParamsAttnMaskType, - kNVTEFusedAttnBwdParamsSoftmaxType, - kNVTEFusedAttnBwdParamsAttnScale, - kNVTEFusedAttnBwdParamsDropout, - kNVTEFusedAttnBwdParamsWindowSizeLeft, - kNVTEFusedAttnBwdParamsWindowSizeRight, - kNVTEFusedAttnBwdParamsBottomRightDiagonal, - kNVTEFusedAttnBwdParamsDeterministic, - kNVTEFusedAttnBwdParamsCudaGraph, + kNVTEFusedAttnBwdParamsMaxSeqlenQ, + kNVTEFusedAttnBwdParamsMaxSeqlenKV, + // workspace and stream kNVTEFusedAttnBwdParamsWorkspace, kNVTEFusedAttnBwdParamsStream, + // number of attributes kNVTEFusedAttnBwdParamsNumAttributes }; -/*! \brief Create a default-initialized fused-attention backward-parameter object. */ +/*! \brief Create a fused-attention backward-parameter object. */ NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void); -/*! \brief Destroy a fused-attention backward-parameter handle. */ +/*! \brief Destroy a fused-attention backward-parameter object. */ void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params); /*! \brief Query an attribute in a fused-attention backward-parameter object. */ @@ -406,33 +413,29 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, const void *buf, size_t size_in_bytes); -/*! \brief Get fused attention backend based on input parameters. +/*! \brief Get fused-attention backend based on input parameters. + * + * This function runs cuDNN frontend's support surface checks, builds cuDNN graphs, + * and caches them if the build is successful. * - * 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. + * \param[in] cfg Fused-attention configuration created by + * ``nvte_create_fused_attn_config()``. + * \param[out] message If cuDNN graphs are built successfully, an empty string; + * if not, a diagnostic message with the reason for rejection. + * Pass NULL to skip diagnostics. + * 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. + * + * \return Fused-attention backend, ``NVTE_F16_arbitrary_seqlen`` or ``NVTE_FP8``, + * if the given configuration is supported; otherwise, ``NVTE_No_Backend``. */ 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. @@ -453,6 +456,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. + * + * \deprecated This function has been deprecated in favor of nvte_get_fused_attn_backend_v2. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, @@ -519,10 +524,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. - */ -void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); - -/*! \brief Compute dot product attention with separate Q, K and V. * * \deprecated This function has been deprecated in favor of nvte_fused_attn_fwd_v2. */ @@ -541,6 +542,16 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); +/*! \brief Compute the backward of the dot product attention with separate Q, K and V. + * + * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_bwd_params()``, + * populate it with ``nvte_set_fused_attn_bwd_params_attribute()`` (or ``FusedAttnBwdParamsWrapper``) setters, and + * destroy it with ``nvte_destroy_fused_attn_bwd_params()``. + * + * \param[in,out] params Opaque fused-attention backward-parameter handle. + */ +void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); + /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * * Notes: @@ -598,10 +609,6 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. - */ -void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); - -/*! \brief Compute the backward of the dot product attention with separate Q, K and V. * * \deprecated This function has been deprecated in favor of nvte_fused_attn_bwd_v2. */ @@ -961,9 +968,9 @@ class AttentionShape { /*! \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. + * It owns an opaque ``NVTEFusedAttnConfig`` handle created by + * ``nvte_create_fused_attn_config()``, and provides a convenient, + * chainable interface for setting every field in ``FusedAttnConfig``. */ class FusedAttnConfigWrapper { public: @@ -1018,38 +1025,28 @@ class FusedAttnConfigWrapper { 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)); + 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_dqkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVLayout, &val, sizeof(val)); + FusedAttnConfigWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasType, &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, + FusedAttnConfigWrapper &set_window_size_left(int64_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigWindowSizeLeft, &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, + FusedAttnConfigWrapper &set_window_size_right(int64_t val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigWindowSizeRight, &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)); + 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_softmax_type(NVTE_Softmax_Type val) noexcept { @@ -1060,52 +1057,54 @@ class FusedAttnConfigWrapper { nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigScalingMode, &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_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)); + FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVDtype, &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)); + FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigODtype, &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)); + FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDODtype, &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)); + FusedAttnConfigWrapper &set_dqkv_dtype(NVTEDType val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVDtype, &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)); + FusedAttnConfigWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVLayout, &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)); + FusedAttnConfigWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigOFormat, &val, sizeof(val)); return *this; } - FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVDtype, &val, sizeof(val)); + FusedAttnConfigWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDOFormat, &val, sizeof(val)); return *this; } - FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigODtype, &val, sizeof(val)); + FusedAttnConfigWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVLayout, &val, sizeof(val)); return *this; } - FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDODtype, &val, sizeof(val)); + 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_dqkv_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVDtype, &val, sizeof(val)); + 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_batch_size(size_t val) noexcept { @@ -1117,7 +1116,7 @@ class FusedAttnConfigWrapper { return *this; } FusedAttnConfigWrapper &set_num_gqa_groups(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumGqaGroups, &val, sizeof(val)); + nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumGQAGroups, &val, sizeof(val)); return *this; } FusedAttnConfigWrapper &set_head_dim_qk(size_t val) noexcept { @@ -1128,6 +1127,22 @@ class FusedAttnConfigWrapper { nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigHeadDimV, &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_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_num_pages_k(size_t val) noexcept { nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumPagesK, &val, sizeof(val)); return *this; @@ -1171,14 +1186,6 @@ class FusedAttnConfigWrapper { 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; - } private: NVTEFusedAttnConfig cfg_ = nullptr; @@ -1186,6 +1193,10 @@ class FusedAttnConfigWrapper { /*! \class FusedAttnFwdParamsWrapper * \brief C++ helper for constructing an ``NVTEFusedAttnFwdParams``. + * + * It owns an opaque ``NVTEFusedAttnFwdParams`` handle created by + * ``nvte_create_fused_attn_fwd_params()``, and provides a convenient, + * chainable interface for setting every field in ``FusedAttnFwdParams``. */ class FusedAttnFwdParamsWrapper { public: @@ -1280,88 +1291,88 @@ class FusedAttnFwdParamsWrapper { sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenQ, &val, - sizeof(val)); + FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsIsTraining, &u8_val, + sizeof(u8_val)); return *this; } - FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenKV, &val, - sizeof(val)); + FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCudaGraph, &u8_val, + sizeof(u8_val)); return *this; } - FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVLayout, &val, - sizeof(val)); + FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsReturnMaxLogit, + &u8_val, sizeof(u8_val)); return *this; } - FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsOFormat, &val, + FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnMaskType, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVScaleInvFormat, - &val, sizeof(val)); - return *this; - } FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBiasType, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnMaskType, &val, + FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxType, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxType, &val, + FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeLeft, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnScale, &val, + FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeRight, &val, sizeof(val)); return *this; } + FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBottomRightDiagonal, + &u8_val, sizeof(u8_val)); + return *this; + } FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsDropout, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeLeft, &val, + FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnScale, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeRight, &val, + FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVLayout, &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBottomRightDiagonal, - &u8_val, sizeof(u8_val)); + FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsOFormat, &val, + sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsIsTraining, &u8_val, - sizeof(u8_val)); + FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + &val, sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsReturnMaxLogit, - &u8_val, sizeof(u8_val)); + FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenQ, &val, + sizeof(val)); return *this; } - FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCudaGraph, &u8_val, - sizeof(u8_val)); + FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenKV, &val, + sizeof(val)); return *this; } FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { @@ -1381,6 +1392,10 @@ class FusedAttnFwdParamsWrapper { /*! \class FusedAttnBwdParamsWrapper * \brief C++ helper for constructing an ``NVTEFusedAttnBwdParams``. + * + * It owns an opaque ``NVTEFusedAttnBwdParams`` handle created by + * ``nvte_create_fused_attn_bwd_params()``, and provides a convenient, + * chainable interface for setting every field in ``FusedAttnBwdParams``. */ class FusedAttnBwdParamsWrapper { public: @@ -1480,97 +1495,97 @@ class FusedAttnBwdParamsWrapper { &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenQ, &val, - sizeof(val)); + FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCudaGraph, &u8_val, + sizeof(u8_val)); return *this; } - FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenKV, &val, - sizeof(val)); + FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDeterministic, &u8_val, + sizeof(u8_val)); return *this; } - FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVLayout, &val, + FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnMaskType, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsOFormat, &val, + FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBiasType, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOFormat, &val, + FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsSoftmaxType, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQKVLayout, &val, + FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeLeft, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, - &val, sizeof(val)); + FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeRight, &val, + sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOScaleInvFormat, &val, - sizeof(val)); + FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + const uint8_t u8_val = static_cast(val); + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBottomRightDiagonal, + &u8_val, sizeof(u8_val)); return *this; } - FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBiasType, &val, + FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDropout, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnMaskType, &val, + FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnScale, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsSoftmaxType, &val, + FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVLayout, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnScale, &val, + FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsOFormat, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDropout, &val, + FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOFormat, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeLeft, &val, + FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQKVLayout, &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeRight, &val, - sizeof(val)); + FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, + &val, sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBottomRightDiagonal, - &u8_val, sizeof(u8_val)); + FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOScaleInvFormat, &val, + sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDeterministic, &u8_val, - sizeof(u8_val)); + FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenQ, &val, + sizeof(val)); return *this; } - FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCudaGraph, &u8_val, - sizeof(u8_val)); + FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenKV, &val, + sizeof(val)); return *this; } FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 6cc1fbb4d5..dc8aca409a 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -130,9 +130,6 @@ class FusedAttnHelper: window_size: Tuple[int, int] bottom_right_diagonal: bool attn_scale: float = 1.0 - # Actual POST_SCALE_BIAS operand dims (may be broadcast, e.g. 1). Left None when the caller does - # not know the bias shape (e.g. the config-level is_fused_attn_kernel_available API), in which - # case get_fused_attn_backend falls back to the full [b, h, sq, skv] representative shape. bias_batch: Optional[int] = None bias_heads: Optional[int] = None bias_seqlen_q: Optional[int] = None @@ -151,27 +148,15 @@ def get_fused_attn_backend(self): """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. + diagnostic string explaining why the configuration was rejected. """ q_type = jax_dtype_to_te_dtype(self.q_dtype) - # The support probe builds a cuDNN graph, which for POST_SCALE_BIAS needs a concrete bias - # shape. Prefer the actual bias operand dims (threaded from the bias aval) so the probe keys - # and builds the exact graph execution uses, even for broadcast bias. When the caller does - # not know the shape (e.g. the config-level is_fused_attn_kernel_available API), fall back to - # the full [b, h, sq, skv] representative shape; backend support does not depend on the bias - # broadcast pattern. For other bias types there is no bias operand, so pass 0 to avoid - # fragmenting the graph-cache key. + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - bias_batch = self.bias_batch if self.bias_batch is not None else self.batch_size - bias_heads = self.bias_heads if self.bias_heads is not None else self.q_num_heads - bias_seqlen_q = ( - self.bias_seqlen_q if self.bias_seqlen_q is not None else self.q_max_seqlen - ) - bias_seqlen_kv = ( - self.bias_seqlen_kv if self.bias_seqlen_kv is not None else self.kv_max_seqlen - ) - else: - bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 + bias_batch = self.bias_batch + bias_heads = self.bias_heads + bias_seqlen_q = self.bias_seqlen_q + bias_seqlen_kv = self.bias_seqlen_kv return transformer_engine_jax.get_fused_attn_backend( self.is_training, self.batch_size, @@ -395,19 +380,10 @@ def abstract( # backend determines the softmax buffer shape/dtype input_batch = reduce(operator.mul, batch_shape) - # Thread the actual POST_SCALE_BIAS operand dims so the trace-time support probe keys and - # builds the exact cuDNN graph the runtime executes (incl. broadcast bias). Derive them the - # same way the lowering/execution does: the bias aval is [*batch, heads, sq, skv] where the - # leading batch dims may be >1 and are collapsed into a single bias_batch. Matching that - # computation keeps the prewarm and runtime graph-cache keys identical for any bias rank. - is_post_scale_bias = config.attn_bias_type == AttnBiasType.POST_SCALE_BIAS - if is_post_scale_bias: - *probe_bias_batch_shape, probe_bias_heads, probe_bias_seqlen_q, probe_bias_seqlen_kv = ( - bias_aval.shape - ) - probe_bias_batch = reduce(operator.mul, probe_bias_batch_shape) - else: - probe_bias_batch = probe_bias_heads = probe_bias_seqlen_q = probe_bias_seqlen_kv = None + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None + if config.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias_aval.shape + bias_batch = reduce(operator.mul, bias_batch_shape) backend, message = FusedAttnHelper( config.is_training, input_batch, @@ -427,10 +403,10 @@ def abstract( config.window_size, config.bottom_right_diagonal, attn_scale=float(config.scaling_factor), - bias_batch=probe_bias_batch, - bias_heads=probe_bias_heads, - bias_seqlen_q=probe_bias_seqlen_q, - bias_seqlen_kv=probe_bias_seqlen_kv, + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ).get_fused_attn_backend() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index dea269a2a8..be5c9a9440 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -232,9 +232,7 @@ class AttentionParams: core_attention_bias_type : str, default = no_bias Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. core_attention_bias_shape : Optional[Tuple[int, int, int, int]], default = None - Broadcast shape of the `core_attention_bias` tensor as `(b, h, sq, skv)`. `None` when no - bias tensor is present. The broadcast pattern (`1hss`, `bhss`, etc.) is derived inside - `get_attention_backend`. + Attention bias shape, (b, h, sq, skv). core_attention_bias_requires_grad : bool, default = True Whether attention bias requires gradient. pad_between_seqs : bool, default = False @@ -267,8 +265,7 @@ class AttentionParams: 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. + Pre-softmax attention scale. fp8_output : bool, default = False Whether output is requested in FP8. checkpoint_core_attention : bool, default = False @@ -312,7 +309,7 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 - softmax_scale: float = 1.0 + softmax_scale: float = 0.0 fp8_output: bool = False checkpoint_core_attention: bool = False has_score_mod: bool = False @@ -356,6 +353,7 @@ class FusedAttentionParams: softmax_type: tex.NVTE_Softmax_Type = tex.NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING dropout: float = 0.0 + attn_scale: float = 1.0 # data types qkv_dtype: DType = DType.kBFloat16 @@ -371,9 +369,6 @@ class FusedAttentionParams: qkv_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET do_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - # attention scaling - attn_scale: float = 0.0 - # tensor dimensions batch_size: int = 0 num_attn_heads: int = 0 diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 84c831f23e..83707a2e16 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -86,10 +86,8 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T * Attention **************************************************************************************************/ -// 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( - py::object fused_attn_params); + const py::object &fused_attn_params); 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 0514f49587..58c35bd8c8 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -40,9 +40,7 @@ void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &s namespace transformer_engine::pytorch { // get the fused attention backend -std::tuple get_fused_attn_backend( - py::object fused_attn_params) { - py::object &p = fused_attn_params; +std::tuple get_fused_attn_backend(const py::object &p) { FusedAttnConfigWrapper cfg; cfg.set_is_training(p.attr("is_training").cast()) .set_deterministic(p.attr("deterministic").cast()) From eadd005c43f007f131105daccd1c42aaa9a7115c Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:18:21 -0700 Subject: [PATCH 30/63] simplify fused attn config/params wrappers via set_attr helper, add graph_debug, fix bias shape handling Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn_score_mod.py | 19 +- tests/pytorch/utils.py | 14 +- .../common/fused_attn/config_and_params.cpp | 3 +- .../common/fused_attn/config_and_params.h | 21 +- .../common/fused_attn/fused_attn.cpp | 5 +- .../fused_attn_f16_arbitrary_seqlen.cu | 35 +- .../common/fused_attn/fused_attn_fp8.cu | 75 +-- .../common/fused_attn/graph_debug.h | 107 ++++ .../include/transformer_engine/fused_attn.h | 463 ++++++------------ .../common/util/pybind_helper.h | 3 +- transformer_engine/jax/attention.py | 19 +- transformer_engine/jax/flax/transformer.py | 23 +- .../attention/dot_product_attention/utils.py | 13 +- .../pytorch/csrc/extensions/attention.cpp | 24 +- 14 files changed, 412 insertions(+), 412 deletions(-) create mode 100644 transformer_engine/common/fused_attn/graph_debug.h diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index b1f165f491..6de133f822 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -18,7 +18,7 @@ ) from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config from transformer_engine.jax.flax import transformer as flax_transformer -from transformer_engine_jax import get_device_compute_capability +from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend from test_fused_attn import FusedAttnRunner, SeqDescFormat @@ -397,9 +397,14 @@ def _identity_score_mod(_graph, score, _tensors): def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True): captured = {} - def fake_fused_attn_kernel_check(*args, **kwargs): - captured.setdefault("kernel_checks", []).append((args, kwargs)) - return kernel_available + class FakeFusedAttnHelper: + def __init__(self, *args, **kwargs): + captured.setdefault("kernel_checks", []).append((args, kwargs)) + + def get_fused_attn_backend(self): + if kernel_available: + return NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, "" + return NVTE_Fused_Attn_Backend.NVTE_No_Backend, "fake: no backend" def fake_fused_attn( qkv, @@ -454,11 +459,7 @@ def fake_fused_attn( ) return qkv[0] - monkeypatch.setattr( - flax_transformer, - "is_fused_attn_kernel_available", - fake_fused_attn_kernel_check, - ) + monkeypatch.setattr(flax_transformer, "FusedAttnHelper", FakeFusedAttnHelper) monkeypatch.setattr(flax_transformer, "fused_attn", fake_fused_attn) return captured diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 89d820ab86..e759ed1c26 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -361,9 +361,15 @@ def get_available_attention_backends( if config.bias_shape == "bhss": alibi_slopes_shape = [config.batch_size, config.num_heads] - core_attention_bias_shape = ( - config.bias_shape if config.attn_bias_type == "post_scale_bias" else None - ) + core_attention_bias_shape = None + if config.attn_bias_type == "post_scale_bias": + b_dim, h_dim, sq_dim, skv_dim = config.bias_shape + core_attention_bias_shape = ( + config.batch_size if b_dim == "b" else 1, + config.num_heads if h_dim == "h" else 1, + config.max_seqlen_q if sq_dim == "s" else 1, + config.max_seqlen_kv if skv_dim == "s" else 1, + ) core_attention_bias_requires_grad = False # d=256 is supported by cuDNN 9.0+ for inference but not training if ( @@ -372,7 +378,7 @@ def get_available_attention_backends( and config.head_dim_v <= 128 ): # TODO(KshitijLakhani): Remove this guard when cuDNN starts support dbias calculation for bias shape 111s - if core_attention_bias_shape != "111s": + if config.bias_shape != "111s": core_attention_bias_requires_grad = True fused_attn_backends = [] diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 33088115af..5433016867 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -78,7 +78,7 @@ void FusedAttnConfig::derive() { } } -FusedAttnConfig FusedAttnConfig::make_cache_key(bool is_forward) const { +FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig cache_cfg = *this; const int64_t s_q = static_cast(cache_cfg.max_seqlen_q); @@ -139,6 +139,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key(bool is_forward) const { FusedAttnConfig FusedAttnFwdParams::make_config() const { const FusedAttnFwdParams ¶ms = *this; FusedAttnConfig cfg{}; + cfg.is_forward = true; cfg.is_training = params.is_training; cfg.deterministic = false; cfg.cuda_graph = params.cuda_graph; diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 1305f21b2f..a68236d79a 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -19,7 +19,7 @@ namespace transformer_engine { struct FusedAttnConfig { - // basic attention knobs + // basic attention settings bool is_training = true; bool deterministic = false; bool cuda_graph = false; @@ -34,13 +34,13 @@ struct FusedAttnConfig { float dropout = 0.0f; float attn_scale = 1.0f; - // data types + // tensor types NVTEDType qkv_dtype = kNVTEBFloat16; NVTEDType o_dtype = kNVTEBFloat16; NVTEDType do_dtype = kNVTEBFloat16; NVTEDType dqkv_dtype = kNVTEBFloat16; - // data and scale layout + // tensor layouts 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; @@ -64,6 +64,10 @@ struct FusedAttnConfig { size_t bucketed_num_tokens_q = 0; size_t bucketed_num_tokens_kv = 0; + // query control (internal only, excluded from attribute serialization, operator<, + // and the graph cache key since it is not a property of the cuDNN graph) + bool is_forward = false; + // paged KV dimensions size_t num_pages_k = 0; size_t num_pages_v = 0; @@ -79,7 +83,7 @@ struct FusedAttnConfig { size_t bias_seqlen_kv = 0; static constexpr size_t attr_sizes[] = { - // basic attention knobs + // basic attention settings sizeof(uint8_t), // is_training sizeof(uint8_t), // deterministic sizeof(uint8_t), // cuda_graph @@ -93,12 +97,12 @@ struct FusedAttnConfig { sizeof(NVTEScalingMode), // scaling_mode sizeof(float), // dropout sizeof(float), // attn_scale - // data types + // tensor types sizeof(NVTEDType), // qkv_dtype sizeof(NVTEDType), // o_dtype sizeof(NVTEDType), // do_dtype sizeof(NVTEDType), // dqkv_dtype - // data and scale layout + // tensor layouts sizeof(NVTE_QKV_Layout), // qkv_layout sizeof(NVTE_QKV_Format), // o_format sizeof(NVTE_QKV_Format), // do_format @@ -162,8 +166,9 @@ struct FusedAttnConfig { // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. // It drops fields that are invariant (e.g. batch_size) or irrelevant (e.g. dO/dQKV dtypes // and `deterministic` for forward, and `return_max_logit` for backward) to the corresponding graph. - // This helps avoid redundant graph builds and cache misses. - FusedAttnConfig make_cache_key(bool is_forward) const; + // This helps avoid redundant graph builds and cache misses. Forward vs. backward is taken from + // the `is_forward` member. + FusedAttnConfig make_cache_key() const; }; inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index c74088713b..c6afbe30f1 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -299,7 +299,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi set_message(message, std::move(fwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg.is_training) { + if (cfg.is_training && !cfg.is_forward) { std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); if (!bwd_reason.empty()) { set_message(message, std::move(bwd_reason)); @@ -324,7 +324,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi set_message(message, std::move(fwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg.is_training) { + if (cfg.is_training && !cfg.is_forward) { std::string bwd_reason = is_supported_f16_bwd(cfg, handle); if (!bwd_reason.empty()) { set_message(message, std::move(bwd_reason)); @@ -353,7 +353,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( cfg.bias_type = bias_type; cfg.attn_mask_type = attn_mask_type; cfg.softmax_type = softmax_type; - cfg.attn_scale = attn_scale; cfg.dropout = dropout; cfg.max_seqlen_q = max_seqlen_q; cfg.max_seqlen_kv = max_seqlen_kv; 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 1fa7c3e9b8..6fd8b60b1d 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 @@ -17,6 +17,7 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" +#include "graph_debug.h" // [GRAPH-DEBUG] #include "utils.h" #define Q_ID 1 @@ -61,12 +62,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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); + const int64_t h = static_cast(cfg.num_attn_heads); + const 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 int64_t d_qk = static_cast(cfg.head_dim_qk); + const 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); @@ -142,7 +143,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; bool generate_stats = true; // Always return stats - const FusedAttnConfig cache_cfg = cfg.make_cache_key(/*is_forward=*/true); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -267,9 +268,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_options.set_diagonal_band_right_bound(window_size_right); - } else if (is_causal || is_bottom_right) { - // Preferred replacement for the deprecated set_causal_mask[_bottom_right]: causal - // masking = diagonal alignment (set above) + a right band bound of 0. + } + if (is_causal || is_bottom_right) { sdpa_options.set_diagonal_band_right_bound(0); } @@ -418,6 +418,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_fwd_build(); // [GRAPH-DEBUG] return return_tuple; }; @@ -448,6 +449,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } + fused_attn_graph_debug::note_fwd_exec(); // [GRAPH-DEBUG] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -557,12 +559,12 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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); + const int64_t h = static_cast(cfg.num_attn_heads); + const 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 int64_t d_qk = static_cast(cfg.head_dim_qk); + const 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); @@ -625,7 +627,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( // 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 = cfg.make_cache_key(/*is_forward=*/false); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); try { namespace fe = cudnn_frontend; @@ -788,9 +790,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } else if (is_causal || is_bottom_right) { - // Preferred replacement for the deprecated set_causal_mask[_bottom_right]: causal - // masking = diagonal alignment (set above) + a right band bound of 0. + } + if (is_causal || is_bottom_right) { sdpa_backward_options.set_diagonal_band_right_bound(0); } @@ -913,6 +914,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_bwd_build(); // [GRAPH-DEBUG] return return_tuple; }; @@ -943,6 +945,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } + fused_attn_graph_debug::note_bwd_exec(); // [GRAPH-DEBUG] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 95f2eaebb4..d4e695473c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -8,6 +8,7 @@ #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" +#include "graph_debug.h" // [GRAPH-DEBUG] #include "utils.h" namespace transformer_engine { @@ -32,13 +33,13 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de 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 int64_t b = static_cast(cfg.batch_size); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t s_q = static_cast(cfg.max_seqlen_q); + const int64_t s_kv = static_cast(cfg.max_seqlen_kv); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const 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; @@ -82,7 +83,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - const FusedAttnConfig cache_cfg = cfg.make_cache_key(/*is_forward=*/true); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -234,10 +235,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de sdpa_options.set_diagonal_band_right_bound(window_size_right); } } - // Preferred replacement for the deprecated set_causal_mask: causal masking = diagonal - // alignment (set above) + a right band bound of 0, unless an explicit right bound was - // already applied above. - if (is_causal && !(cudnn_runtime_version >= 92100 && window_size_right != -1)) { + if (is_causal) { sdpa_options.set_diagonal_band_right_bound(0); } @@ -358,6 +356,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_fwd_build(); // [GRAPH-DEBUG] return return_tuple; }; @@ -374,6 +373,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } + fused_attn_graph_debug::note_fwd_exec(); // [GRAPH-DEBUG] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -459,13 +459,13 @@ void fused_attn_fp8_bwd_impl( 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 int64_t b = static_cast(cfg.batch_size); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t s_q = static_cast(cfg.max_seqlen_q); + const int64_t s_kv = static_cast(cfg.max_seqlen_kv); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const 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; @@ -515,7 +515,7 @@ 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 = cfg.make_cache_key(/*is_forward=*/false); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); try { namespace fe = cudnn_frontend; using graph_and_tensors = @@ -786,10 +786,7 @@ void fused_attn_fp8_bwd_impl( sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); } } - // Preferred replacement for the deprecated set_causal_mask: causal masking = diagonal - // alignment (set above) + a right band bound of 0, unless an explicit right bound was - // already applied above. - if (is_causal && !(cudnn_runtime_version >= 92100 && window_size_right != -1)) { + if (is_causal) { sdpa_backward_options.set_diagonal_band_right_bound(0); } @@ -962,6 +959,7 @@ void fused_attn_fp8_bwd_impl( std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_bwd_build(); // [GRAPH-DEBUG] return return_tuple; }; @@ -979,6 +977,7 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } + fused_attn_graph_debug::note_bwd_exec(); // [GRAPH-DEBUG] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1153,11 +1152,14 @@ void fused_attn_fp8_fwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const size_t workspace_size = 0; + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); + 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( - cfg, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, + graph_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); @@ -1265,15 +1267,18 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const size_t workspace_size = 0; + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); + 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( - 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, + graph_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, workspace->data.dptr, &workspace_size, stream, handle); @@ -1295,10 +1300,13 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const } std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); + size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_fwd_impl( - cfg, + graph_cfg, /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, /*devPtrDescaleQ=*/nullptr, /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, @@ -1317,10 +1325,13 @@ std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handl } std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); + size_t workspace_size = 0; try { fused_attn::fused_attn_fp8_bwd_impl( - cfg, + graph_cfg, /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, diff --git a/transformer_engine/common/fused_attn/graph_debug.h b/transformer_engine/common/fused_attn/graph_debug.h new file mode 100644 index 0000000000..c5afc8f035 --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_debug.h @@ -0,0 +1,107 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// ============================================================================ +// [GRAPH-DEBUG] TEMPORARY DEBUG INSTRUMENTATION -- REMOVE AFTER VERIFICATION. +// +// Counts fused-attention cuDNN graph *builds* (cache misses that construct a new +// graph) vs. *executions* (real forward/backward runs, excluding workspace-sizing +// probes) to detect redundant graph construction. +// +// Enable at runtime with: export NVTE_FUSED_ATTN_GRAPH_DEBUG=1 +// A running "BUILD" line is printed whenever a new graph is constructed, and a +// "SUMMARY" line with final totals is printed at process exit. +// +// To remove all of this instrumentation later: +// 1. Delete this file (graph_debug.h). +// 2. Remove every line tagged with the "[GRAPH-DEBUG]" marker in: +// - fused_attn_fp8.cu +// - fused_attn_f16_arbitrary_seqlen.cu +// ============================================================================ + +#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ +#define TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ + +#include +#include +#include +#include + +namespace transformer_engine { +namespace fused_attn_graph_debug { + +inline std::atomic &fwd_built() { + static std::atomic v{0}; + return v; +} +inline std::atomic &fwd_exec() { + static std::atomic v{0}; + return v; +} +inline std::atomic &bwd_built() { + static std::atomic v{0}; + return v; +} +inline std::atomic &bwd_exec() { + static std::atomic v{0}; + return v; +} + +inline bool enabled() { + static const bool on = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return on; +} + +inline void dump(const char *event) { + std::fprintf(stderr, + "[GRAPH-DEBUG] %-10s | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", event, + static_cast(fwd_built().load()), + static_cast(fwd_exec().load()), + static_cast(bwd_built().load()), + static_cast(bwd_exec().load())); + std::fflush(stderr); +} + +inline void register_summary_once() { + static const bool registered = [] { + std::atexit([] { + if (enabled()) dump("SUMMARY"); + }); + return true; + }(); + (void)registered; +} + +inline void note_fwd_build() { + if (!enabled()) return; + register_summary_once(); + fwd_built().fetch_add(1); + dump("fwd BUILD"); +} +inline void note_fwd_exec() { + if (!enabled()) return; + register_summary_once(); + fwd_exec().fetch_add(1); +} +inline void note_bwd_build() { + if (!enabled()) return; + register_summary_once(); + bwd_built().fetch_add(1); + dump("bwd BUILD"); +} +inline void note_bwd_exec() { + if (!enabled()) return; + register_summary_once(); + bwd_exec().fetch_add(1); +} + +} // namespace fused_attn_graph_debug +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 3010e6418a..0b2dfd6b85 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -207,7 +207,7 @@ typedef void *NVTEFusedAttnConfig; * at the end and existing fields are never reordered, removed, or resized. */ enum NVTEFusedAttnConfigAttribute { - // basic configuration knobs + // basic attention settings kNVTEFusedAttnConfigIsTraining = 0, kNVTEFusedAttnConfigDeterministic, kNVTEFusedAttnConfigCudaGraph, @@ -542,6 +542,16 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); +/*! \brief Compute dot product attention with separate Q, K and V. + * + * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_fwd_params()``, + * populate it with ``nvte_set_fused_attn_fwd_params_attribute()`` (or ``FusedAttnFwdParamsWrapper``) setters, and + * destroy it with ``nvte_destroy_fused_attn_fwd_params()``. + * + * \param[in,out] params Opaque fused-attention forward-parameter handle. + */ +void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); + /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_bwd_params()``, @@ -1002,192 +1012,141 @@ class FusedAttnConfigWrapper { 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; + return set_attr(kNVTEFusedAttnConfigIsTraining, static_cast(val)); } 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; + return set_attr(kNVTEFusedAttnConfigDeterministic, static_cast(val)); } 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; + return set_attr(kNVTEFusedAttnConfigCudaGraph, static_cast(val)); } 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; + return set_attr(kNVTEFusedAttnConfigReturnMaxLogit, static_cast(val)); } FusedAttnConfigWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigAttnMaskType, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigAttnMaskType, val); } FusedAttnConfigWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasType, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigBiasType, val); } FusedAttnConfigWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigWindowSizeLeft, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigWindowSizeLeft, val); } FusedAttnConfigWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigWindowSizeRight, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigWindowSizeRight, val); } 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; + return set_attr(kNVTEFusedAttnConfigBottomRightDiagonal, static_cast(val)); } FusedAttnConfigWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigSoftmaxType, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigSoftmaxType, val); } FusedAttnConfigWrapper &set_scaling_mode(NVTEScalingMode val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigScalingMode, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigScalingMode, val); } FusedAttnConfigWrapper &set_dropout(float val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDropout, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigDropout, val); } FusedAttnConfigWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigAttnScale, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigAttnScale, val); } FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVDtype, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigQKVDtype, val); } FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigODtype, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigODtype, val); } FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDODtype, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigDODtype, val); } FusedAttnConfigWrapper &set_dqkv_dtype(NVTEDType val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVDtype, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigDQKVDtype, val); } FusedAttnConfigWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVLayout, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigQKVLayout, val); } FusedAttnConfigWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigOFormat, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigOFormat, val); } FusedAttnConfigWrapper &set_do_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDOFormat, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigDOFormat, val); } FusedAttnConfigWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDQKVLayout, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigDQKVLayout, val); } FusedAttnConfigWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigQKVScaleInvFormat, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigQKVScaleInvFormat, val); } FusedAttnConfigWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigDOScaleInvFormat, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigDOScaleInvFormat, val); } FusedAttnConfigWrapper &set_batch_size(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBatchSize, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigBatchSize, val); } FusedAttnConfigWrapper &set_num_attn_heads(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumAttnHeads, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigNumAttnHeads, val); } FusedAttnConfigWrapper &set_num_gqa_groups(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumGQAGroups, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigNumGQAGroups, val); } FusedAttnConfigWrapper &set_head_dim_qk(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigHeadDimQK, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigHeadDimQK, val); } FusedAttnConfigWrapper &set_head_dim_v(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigHeadDimV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigHeadDimV, val); } FusedAttnConfigWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxSeqlenQ, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigMaxSeqlenQ, val); } FusedAttnConfigWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxSeqlenKV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigMaxSeqlenKV, val); } FusedAttnConfigWrapper &set_num_tokens_q(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumTokensQ, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigNumTokensQ, val); } FusedAttnConfigWrapper &set_num_tokens_kv(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumTokensKV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigNumTokensKV, val); } FusedAttnConfigWrapper &set_num_pages_k(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumPagesK, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigNumPagesK, val); } FusedAttnConfigWrapper &set_num_pages_v(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigNumPagesV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigNumPagesV, val); } FusedAttnConfigWrapper &set_page_size_k(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigPageSizeK, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigPageSizeK, val); } FusedAttnConfigWrapper &set_page_size_v(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigPageSizeV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigPageSizeV, val); } FusedAttnConfigWrapper &set_max_pages_per_seq_k(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxPagesPerSeqK, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigMaxPagesPerSeqK, val); } FusedAttnConfigWrapper &set_max_pages_per_seq_v(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigMaxPagesPerSeqV, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigMaxPagesPerSeqV, val); } FusedAttnConfigWrapper &set_bias_batch_size(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasBatchSize, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigBiasBatchSize, val); } FusedAttnConfigWrapper &set_bias_num_heads(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasNumHeads, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigBiasNumHeads, val); } FusedAttnConfigWrapper &set_bias_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasSeqlenQ, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigBiasSeqlenQ, val); } FusedAttnConfigWrapper &set_bias_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, kNVTEFusedAttnConfigBiasSeqlenKV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnConfigBiasSeqlenKV, val); } private: + // Common implementation for every setter: copy the value to a local, forward + // its address and size to the C API, and return *this for chaining. + template + FusedAttnConfigWrapper &set_attr(NVTEFusedAttnConfigAttribute attr, T val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, attr, &val, sizeof(val)); + return *this; + } + NVTEFusedAttnConfig cfg_ = nullptr; }; @@ -1201,11 +1160,14 @@ class FusedAttnConfigWrapper { class FusedAttnFwdParamsWrapper { public: FusedAttnFwdParamsWrapper() : params_{nvte_create_fused_attn_fwd_params()} {} + FusedAttnFwdParamsWrapper(const FusedAttnFwdParamsWrapper &) = delete; FusedAttnFwdParamsWrapper &operator=(const FusedAttnFwdParamsWrapper &) = delete; + FusedAttnFwdParamsWrapper(FusedAttnFwdParamsWrapper &&other) noexcept : params_{other.params_} { other.params_ = nullptr; } + FusedAttnFwdParamsWrapper &operator=(FusedAttnFwdParamsWrapper &&other) noexcept { if (this != &other) { nvte_destroy_fused_attn_fwd_params(params_); @@ -1214,179 +1176,125 @@ class FusedAttnFwdParamsWrapper { } return *this; } + ~FusedAttnFwdParamsWrapper() { if (params_ != nullptr) { nvte_destroy_fused_attn_fwd_params(params_); } } + operator NVTEFusedAttnFwdParams() const noexcept { return params_; } NVTEFusedAttnFwdParams get() const noexcept { return params_; } + FusedAttnFwdParamsWrapper &set_Q(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQ, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsQ, val); } FusedAttnFwdParamsWrapper &set_K(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsK, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsK, val); } FusedAttnFwdParamsWrapper &set_V(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsV, val); } FusedAttnFwdParamsWrapper &set_Bias(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBias, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsBias, val); } FusedAttnFwdParamsWrapper &set_SoftmaxOffset(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxOffset, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsSoftmaxOffset, val); } FusedAttnFwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQ, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQ, val); } FusedAttnFwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKV, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensKV, val); } FusedAttnFwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensQPadded, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQPadded, val); } FusedAttnFwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, - &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, val); } FusedAttnFwdParamsWrapper &set_page_table_k(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableK, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsPageTableK, val); } FusedAttnFwdParamsWrapper &set_page_table_v(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsPageTableV, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsPageTableV, val); } FusedAttnFwdParamsWrapper &set_rng_state(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsRngState, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsRngState, val); } FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsS, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsS, val); } FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsO, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsO, val); } FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack *val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAuxCtxTensors, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsAuxCtxTensors, val); } FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsIsTraining, &u8_val, - sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsIsTraining, static_cast(val)); } FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsCudaGraph, &u8_val, - sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsCudaGraph, static_cast(val)); } FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsReturnMaxLogit, - &u8_val, sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); } FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnMaskType, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); } FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBiasType, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsBiasType, val); } FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsSoftmaxType, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsSoftmaxType, val); } FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeLeft, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsWindowSizeLeft, val); } FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWindowSizeRight, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsWindowSizeRight, val); } FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsBottomRightDiagonal, - &u8_val, sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsBottomRightDiagonal, static_cast(val)); } FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsDropout, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsDropout, val); } FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsAttnScale, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); } FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVLayout, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsQKVLayout, val); } FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsOFormat, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsOFormat, val); } FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsQKVScaleInvFormat, - &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsQKVScaleInvFormat, val); } FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenQ, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenQ, val); } FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsMaxSeqlenKV, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenKV, val); } FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsWorkspace, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsWorkspace, val); } FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, kNVTEFusedAttnFwdParamsStream, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnFwdParamsStream, val); } private: + // Common implementation for every setter: copy the value to a local, forward + // its address and size to the C API, and return *this for chaining. + template + FusedAttnFwdParamsWrapper &set_attr(NVTEFusedAttnFwdParamsAttribute attr, T val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, attr, &val, sizeof(val)); + return *this; + } + NVTEFusedAttnFwdParams params_ = nullptr; }; @@ -1400,11 +1308,14 @@ class FusedAttnFwdParamsWrapper { class FusedAttnBwdParamsWrapper { public: FusedAttnBwdParamsWrapper() : params_{nvte_create_fused_attn_bwd_params()} {} + FusedAttnBwdParamsWrapper(const FusedAttnBwdParamsWrapper &) = delete; FusedAttnBwdParamsWrapper &operator=(const FusedAttnBwdParamsWrapper &) = delete; + FusedAttnBwdParamsWrapper(FusedAttnBwdParamsWrapper &&other) noexcept : params_{other.params_} { other.params_ = nullptr; } + FusedAttnBwdParamsWrapper &operator=(FusedAttnBwdParamsWrapper &&other) noexcept { if (this != &other) { nvte_destroy_fused_attn_bwd_params(params_); @@ -1413,193 +1324,137 @@ class FusedAttnBwdParamsWrapper { } return *this; } + ~FusedAttnBwdParamsWrapper() { if (params_ != nullptr) { nvte_destroy_fused_attn_bwd_params(params_); } } + operator NVTEFusedAttnBwdParams() const noexcept { return params_; } NVTEFusedAttnBwdParams get() const noexcept { return params_; } + FusedAttnBwdParamsWrapper &set_Q(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQ, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsQ, val); } FusedAttnBwdParamsWrapper &set_K(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsK, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsK, val); } FusedAttnBwdParamsWrapper &set_V(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsV, val); } FusedAttnBwdParamsWrapper &set_O(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsO, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsO, val); } FusedAttnBwdParamsWrapper &set_dO(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDO, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDO, val); } FusedAttnBwdParamsWrapper &set_S(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsS, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsS, val); } FusedAttnBwdParamsWrapper &set_dP(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDP, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDP, val); } FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(const NVTETensorPack *val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAuxCtxTensors, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsAuxCtxTensors, val); } FusedAttnBwdParamsWrapper &set_dQ(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQ, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDQ, val); } FusedAttnBwdParamsWrapper &set_dK(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDK, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDK, val); } FusedAttnBwdParamsWrapper &set_dV(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDV, &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDV, val); } FusedAttnBwdParamsWrapper &set_dBias(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDBias, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDBias, val); } FusedAttnBwdParamsWrapper &set_dSoftmaxOffset(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDSoftmaxOffset, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDSoftmaxOffset, val); } FusedAttnBwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQ, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensQ, val); } FusedAttnBwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKV, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKV, val); } FusedAttnBwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensQPadded, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensQPadded, val); } FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, - &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, val); } FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsCudaGraph, &u8_val, - sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsCudaGraph, static_cast(val)); } FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDeterministic, &u8_val, - sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDeterministic, static_cast(val)); } FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnMaskType, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); } FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBiasType, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); } FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsSoftmaxType, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsSoftmaxType, val); } FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeLeft, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsWindowSizeLeft, val); } FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWindowSizeRight, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsWindowSizeRight, val); } FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - const uint8_t u8_val = static_cast(val); - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsBottomRightDiagonal, - &u8_val, sizeof(u8_val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsBottomRightDiagonal, static_cast(val)); } FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDropout, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDropout, val); } FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsAttnScale, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); } FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVLayout, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsQKVLayout, val); } FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsOFormat, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsOFormat, val); } FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOFormat, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDOFormat, val); } FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDQKVLayout, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDQKVLayout, val); } FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, - &val, sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsQKVScaleInvFormat, val); } FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsDOScaleInvFormat, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsDOScaleInvFormat, val); } FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenQ, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenQ, val); } FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsMaxSeqlenKV, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenKV, val); } FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsWorkspace, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsWorkspace, val); } FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, kNVTEFusedAttnBwdParamsStream, &val, - sizeof(val)); - return *this; + return set_attr(kNVTEFusedAttnBwdParamsStream, val); } private: + // Common implementation for every setter: copy the value to a local, forward + // its address and size to the C API, and return *this for chaining. + template + FusedAttnBwdParamsWrapper &set_attr(NVTEFusedAttnBwdParamsAttribute attr, T val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, attr, &val, sizeof(val)); + return *this; + } + NVTEFusedAttnBwdParams params_ = nullptr; }; #endif // __cplusplus diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index 540b38143d..d739965163 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -86,7 +86,8 @@ .value("NVTE_Paged_KV_SBHD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD) \ .value("NVTE_Paged_KV_THD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD) \ .value("NVTE_Paged_KV_THD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD) \ - .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD); \ + .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD) \ + .value("NVTE_QKV_Layout_NOT_SET", NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET); \ pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) \ .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) \ .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8) \ diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 3cb3a1dd26..744e81d7d8 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -13,7 +13,6 @@ 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 @@ -342,13 +341,17 @@ def is_fused_attn_kernel_available( head_dim_v, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, - return_reason: bool = False, + bias_batch: Optional[int] = None, + bias_heads: Optional[int] = None, + bias_seqlen_q: Optional[int] = None, + bias_seqlen_kv: Optional[int] = None, ): """ 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). + For a ``POST_SCALE_BIAS`` config, pass the bias broadcast shape via ``bias_batch``, + ``bias_heads``, ``bias_seqlen_q``, and ``bias_seqlen_kv`` so the backend probe matches the + graph used at execution time. """ window_size_tuple = (-1, -1) if window_size is None else window_size @@ -376,13 +379,13 @@ def make_helper(attn_mask_type): head_dim_v, window_size_tuple, bottom_right, + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ) 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() diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 6b8107ee67..276e476aeb 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -5,6 +5,7 @@ Wrapper module for Transformer related layers with FP8 support. """ import functools +import operator from enum import Enum from math import sqrt import os @@ -20,6 +21,7 @@ from jax import random as jax_random from jax import lax, vmap from jax.ad_checkpoint import checkpoint_name +from transformer_engine_jax import NVTE_Fused_Attn_Backend from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP from .module import LayerNorm, Softmax @@ -30,9 +32,10 @@ QKVLayout, SequenceDescriptor, ) -from ..attention import is_fused_attn_kernel_available, make_swa_mask, canonicalize_attn_mask_type +from ..attention import make_swa_mask, canonicalize_attn_mask_type from ..attention import fused_attn from ..attention import CPStrategy +from ..cpp_extensions import FusedAttnHelper from ..softmax import SoftmaxFusionType from ..sharding import num_of_devices from ..sharding import get_sharding_map_logic_axis_to_mesh_axis @@ -797,7 +800,13 @@ 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, fused_attn_reject_reason = is_fused_attn_kernel_available( + # Thread the POST_SCALE_BIAS broadcast shape through so this pre-check probes the same + # cuDNN graph as the primitive does at trace time (see FusedAttnFwdPrimitive.abstract). + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None + if attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias.shape + bias_batch = functools.reduce(operator.mul, bias_batch_shape) + fused_attn_helper = FusedAttnHelper( # This needs to be fixed: TE-Jax has historically correlated training mode # with deterministic mode. not deterministic, @@ -817,9 +826,15 @@ def __call__( seqlen_kv, head_dim_qk, head_dim_v, - self.window_size, - return_reason=True, + (-1, -1) if self.window_size is None else self.window_size, + attn_mask_type.is_bottom_right(), + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ) + fused_attn_backend, fused_attn_reject_reason = fused_attn_helper.get_fused_attn_backend() + has_fused_attn_kernel = fused_attn_backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index be5c9a9440..ed1137a077 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -340,7 +340,7 @@ class FusedAttentionParams: Attention parameters used by the `FusedAttention` backend. """ - # basic attention knobs + # basic attention settings is_training: bool = True deterministic: bool = False cuda_graph: bool = False @@ -355,13 +355,13 @@ class FusedAttentionParams: dropout: float = 0.0 attn_scale: float = 1.0 - # data types + # tensor types qkv_dtype: DType = DType.kBFloat16 o_dtype: DType = DType.kBFloat16 do_dtype: DType = DType.kBFloat16 dqkv_dtype: DType = DType.kBFloat16 - # data and scale layout + # tensor layouts qkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET o_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET do_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET @@ -1429,8 +1429,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fu_core_attention_bias_shape_type = "111s" if sq == 1 and max_seqlen_q != 1 else "11ss" else: raise ValueError( - f"core_attention_bias tensor must be in one of " - "{"bhss", "1hss", "b1ss", "11ss", "111s"} shapes. Found (b,h,sq,skv) = ({b},{h},{sq},{_skv})" + "core_attention_bias tensor must be in one of " + f'{{"bhss", "1hss", "b1ss", "11ss", "111s"}} shapes. ' + f"Found (b,h,sq,skv) = ({b},{h},{sq},{_skv})" ) if ( use_fused_attention @@ -1503,6 +1504,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt softmax_type=SoftmaxType[softmax_type], scaling_mode=scaling_mode, dropout=attention_dropout, + attn_scale=softmax_scale, qkv_dtype=qkv_type, o_dtype=o_type, do_dtype=do_type, @@ -1513,7 +1515,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt dqkv_layout=QKVLayout[dqkv_layout], qkv_scale_inv_format=QKVFormat[qkv_scale_inv_format], do_scale_inv_format=QKVFormat[do_scale_inv_format], - attn_scale=softmax_scale, batch_size=batch_size, num_attn_heads=num_heads, num_gqa_groups=num_gqa_groups, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 58c35bd8c8..5bfbc79d53 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -313,18 +313,17 @@ std::vector fused_attn_fwd( .set_window_size_left(window_size[0]) .set_window_size_right(window_size[1]) .set_bottom_right_diagonal(bottom_right_diagonal) + .set_workspace(workspace.data()) .set_stream(at::cuda::getCurrentCUDAStream()); // populate tensors with appropriate shapes and dtypes - NVTE_SCOPED_GIL_RELEASE({ - params.set_workspace(workspace.data()); - nvte_fused_attn_fwd_v2(params); - }); + NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd_v2(params); }); // allocate memory for workspace and auxiliary output tensors auto workspace_data = allocateSpace(workspace.shape(), workspace.dtype()); workspace = makeTransformerEngineTensor(workspace_data.data_ptr(), workspace.shape(), workspace.dtype()); + params.set_workspace(workspace.data()); // output_tensors = [O, nvte_aux_tensor_pack.tensors] std::vector output_tensors; @@ -367,10 +366,7 @@ std::vector fused_attn_fwd( } // execute the kernel - NVTE_SCOPED_GIL_RELEASE({ - params.set_workspace(workspace.data()); - nvte_fused_attn_fwd_v2(params); - }); + NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd_v2(params); }); // destroy tensor wrappers, but not allocated memory nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); @@ -668,24 +664,20 @@ std::vector fused_attn_bwd( .set_bottom_right_diagonal(bottom_right_diagonal) .set_deterministic(deterministic) .set_cuda_graph(cuda_graph) + .set_workspace(workspace.data()) .set_stream(at::cuda::getCurrentCUDAStream()); // populate tensors with appropriate shapes and dtypes - NVTE_SCOPED_GIL_RELEASE({ - params.set_workspace(workspace.data()); - nvte_fused_attn_bwd_v2(params); - }); + NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); // allocate memory for workspace auto workspace_data = allocateSpace(workspace.shape(), workspace.dtype()); workspace = makeTransformerEngineTensor(workspace_data.data_ptr(), workspace.shape(), workspace.dtype()); + params.set_workspace(workspace.data()); // execute kernel - NVTE_SCOPED_GIL_RELEASE({ - params.set_workspace(workspace.data()); - nvte_fused_attn_bwd_v2(params); - }); + NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); // destroy tensor wrappers nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); From 28f5a8c1ef11fc66c5f278304e0eb2a4458a5f92 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 17 Jul 2026 09:35:45 -0700 Subject: [PATCH 31/63] clean up derived fields, debug probe/exec graph mismatches Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/utils.py | 6 +- .../common/fused_attn/config_and_params.cpp | 107 +++++++------ .../common/fused_attn/config_and_params.h | 42 +++-- .../fused_attn_f16_arbitrary_seqlen.cu | 86 +++++----- .../common/fused_attn/fused_attn_fp8.cu | 36 +++-- .../common/fused_attn/graph_debug.h | 149 +++++++++++++++++- 6 files changed, 295 insertions(+), 131 deletions(-) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index e759ed1c26..4aed95cb2c 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -311,9 +311,9 @@ 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.bottom_right_diagonal = self.attn_mask_type not in { + "causal", + "padding_causal", } self.context_parallel = context_parallel self.cp_comm_type = cp_comm_type diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 5433016867..40b7831ace 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -7,6 +7,7 @@ #include "config_and_params.h" #include +#include #include @@ -39,23 +40,37 @@ void FusedAttnConfig::derive() { const int64_t sq = static_cast(max_seqlen_q); const int64_t skv = static_cast(max_seqlen_kv); - const NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - const NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + // convenience fields + q_format = nvte_get_q_format(qkv_layout); + kv_format = nvte_get_kv_format(qkv_layout); const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - const bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + is_padding = (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); + is_causal = (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK); + is_causal_bottom_right = (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + // bucket the THD (ragged) batch and token counts const size_t tokens_q = num_tokens_q != 0 ? num_tokens_q : static_cast(b * sq); const size_t tokens_kv = num_tokens_kv != 0 ? 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); bucketed_batch_size = (is_ragged_q || is_ragged_kv) ? fused_attn::get_max_batch_size(batch_size) : 0; bucketed_num_tokens_q = is_ragged_q ? fused_attn::get_max_tokens(tokens_q) : 0; bucketed_num_tokens_kv = is_ragged_kv ? fused_attn::get_max_tokens(tokens_kv) : 0; + // use of cu_seqlens vs actual_seqlens + const size_t cudnn_runtime_version = cudnnGetVersion(); + const bool is_dropout = is_training && dropout != 0.0f; + uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12500 && + (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && + !is_dropout; + + // paged KV dimensions if (is_paged_kv) { if (num_pages_k == 0) { num_pages_k = static_cast(b); @@ -81,54 +96,52 @@ void FusedAttnConfig::derive() { FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig cache_cfg = *this; - 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) { + // Normalize bottom_right_diagonal (the cuDNN diagonal alignment). The impl only turns it into a + // real causal band under `is_causal || is_causal_bottom_right` or a sliding window; otherwise the + // alignment is inert, so canonicalize it (like attn_scale) to false. This keeps the backend + // support probe (which passes a possibly-different brd, e.g. default false) and the real op on a + // single cached graph. + const bool has_window = cache_cfg.window_size_left != -1 || cache_cfg.window_size_right != -1; + if (!cache_cfg.is_causal && !cache_cfg.is_causal_bottom_right && !has_window) { + cache_cfg.bottom_right_diagonal = false; + } else if (cache_cfg.is_causal_bottom_right && cache_cfg.max_seqlen_q == cache_cfg.max_seqlen_kv && + !cache_cfg.is_padding) { + // square bottom-right causal collapses to top-left causal (mirrors the impl). 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; + // Bucket THD (ragged) batch and token counts + if (cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) { + const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); + if (cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + if (cache_cfg.is_ragged_q) { + cache_cfg.max_seqlen_q = cache_cfg.bucketed_num_tokens_q; + } + if (cache_cfg.is_ragged_kv) { + cache_cfg.max_seqlen_kv = cache_cfg.bucketed_num_tokens_kv; + } + cache_cfg.num_tokens_q = 0; + cache_cfg.num_tokens_kv = 0; + const bool bucket_batch = !is_forward || !cache_cfg.uses_cu_seqlens_directly; + if (bucket_batch) { + cache_cfg.batch_size = cache_cfg.bucketed_batch_size; + } } } - // cuDNN graph supports dynamic shapes for batch_size - cache_cfg.batch_size = 1; - cache_cfg.bucketed_batch_size = 1; + // attn_scale is a pass-by-value graph input and different scales can share the same cached graph cache_cfg.attn_scale = 1.0f; - // Drop from each graph's cache key the fields the graph does not actually consume, so a graph - // prewarmed by a backend probe (which may carry different values for those ignored fields, e.g. - // the framework get_attention_backend probe) is still reused at execution. The forward graph - // produces O (and optionally softmax stats / max logit) but never consumes the dO/dQKV dtypes or - // the backward-only determinism choice. The backward graph consumes dO/dQKV and honors - // determinism but never produces the forward max-logit output. + // Restrict each direction's key to the fields its graph actually consumes, so + // no redundant graphs are built and no cache misses either if (is_forward) { - if (cache_cfg.is_training) { - cache_cfg.do_dtype = kNVTEBFloat16; - cache_cfg.dqkv_dtype = kNVTEBFloat16; - cache_cfg.deterministic = false; - } + cache_cfg.do_dtype = kNVTEBFloat16; + cache_cfg.dqkv_dtype = kNVTEBFloat16; + cache_cfg.do_format = NVTE_QKV_Format_NOT_SET; + cache_cfg.dqkv_layout = NVTE_QKV_Layout_NOT_SET; + cache_cfg.do_scale_inv_format = NVTE_QKV_Format_NOT_SET; + cache_cfg.deterministic = false; } else { cache_cfg.return_max_logit = false; } diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index a68236d79a..5ec6b02822 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -59,15 +59,6 @@ struct FusedAttnConfig { size_t num_tokens_q = 0; size_t num_tokens_kv = 0; - // derived tensor dimensions (internal only) - size_t bucketed_batch_size = 0; - size_t bucketed_num_tokens_q = 0; - size_t bucketed_num_tokens_kv = 0; - - // query control (internal only, excluded from attribute serialization, operator<, - // and the graph cache key since it is not a property of the cuDNN graph) - bool is_forward = false; - // paged KV dimensions size_t num_pages_k = 0; size_t num_pages_v = 0; @@ -82,6 +73,28 @@ struct FusedAttnConfig { size_t bias_seqlen_q = 0; size_t bias_seqlen_kv = 0; + // Internal-only fields: never part of attribute serialization, operator<, or the graph cache key. + // Filled by derive() or set by caller (i.e. is_forward). Added for convinence purposes and do not + // represent graph properties. + + // Direction to build the cuDNN graph for; steers make_cache_key() normalization. + bool is_forward = false; + // THD batch/token counts; make_cache_key() folds these into batch_size/max_seqlen_*. + size_t bucketed_batch_size = 0; + size_t bucketed_num_tokens_q = 0; + size_t bucketed_num_tokens_kv = 0; + // Uses cu_seqlens or actual_seqlens. + bool uses_cu_seqlens_directly = false; + // Convinence fields to avoid recompute. + NVTE_QKV_Format q_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format kv_format = NVTE_QKV_Format_NOT_SET; + bool is_ragged_q = false; + bool is_ragged_kv = false; + bool is_paged_kv = false; + bool is_padding = false; + bool is_causal = false; + bool is_causal_bottom_right = false; + static constexpr size_t attr_sizes[] = { // basic attention settings sizeof(uint8_t), // is_training @@ -140,8 +153,7 @@ struct FusedAttnConfig { dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, - num_tokens_kv, bucketed_batch_size, bucketed_num_tokens_q, - bucketed_num_tokens_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, + num_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_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv) < std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, @@ -152,8 +164,7 @@ struct FusedAttnConfig { rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, rhs.num_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, rhs.max_seqlen_q, - rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, rhs.bucketed_batch_size, - rhs.bucketed_num_tokens_q, rhs.bucketed_num_tokens_kv, rhs.num_pages_k, + rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, 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); @@ -164,10 +175,9 @@ struct FusedAttnConfig { void derive(); // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. - // It drops fields that are invariant (e.g. batch_size) or irrelevant (e.g. dO/dQKV dtypes + // It drops fields that are invariant (e.g. attn_scale) or irrelevant (e.g. dO/dQKV dtypes // and `deterministic` for forward, and `return_max_logit` for backward) to the corresponding graph. - // This helps avoid redundant graph builds and cache misses. Forward vs. backward is taken from - // the `is_forward` member. + // This helps avoid redundant graph builds and cache misses. FusedAttnConfig make_cache_key() const; }; 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 623d0cb35d..09fae572fb 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 @@ -83,7 +83,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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; + 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; @@ -98,29 +98,21 @@ void fused_attn_arbitrary_seqlen_fwd_impl( bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); - bool is_bottom_right = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - if (is_bottom_right && s_q == s_kv && !is_padding) { - is_causal = true; - is_bottom_right = false; - bottom_right_diagonal = false; - } + bool is_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (is_training && dropout_probability != 0.0f); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + NVTE_QKV_Format q_format = cfg.q_format; + NVTE_QKV_Format kv_format = cfg.kv_format; + bool is_ragged_q = cfg.is_ragged_q; + bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + bool is_paged_kv = cfg.is_paged_kv; if (is_paged_kv) { NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } @@ -128,16 +120,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. - const bool use_cu_seqlens_directly = - CUDNN_FRONTEND_VERSION >= 12500 && - // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN - // version, so we'll do the same. - (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && - // This extra restriction is needed because cuDNN frontend doesn't yet allow - // the combination of dropout and stats generation for the fprop unified engine, - // so any such request would always get routed to the old composite SDPA engine - // (which doesn't support cu_seqlens). Remove this restriction when possible. - !is_dropout; + // Defined on FusedAttnConfig so make_cache_key() keys the graph on the matching batch + // handling (real batch here, bucketed batch on the legacy path); keep the two in sync. + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; @@ -203,7 +188,15 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - if (it != cache.end()) { + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] + sm_arch_ != 120) { // [GRAPH-DEBUG] + fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] + "fwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), + /*legacy=*/!use_cu_seqlens_directly); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } @@ -303,7 +296,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_options.set_diagonal_band_right_bound(window_size_right); } - if (is_causal || is_bottom_right) { + if (is_causal || is_causal_bottom_right) { sdpa_options.set_diagonal_band_right_bound(0); } @@ -660,7 +653,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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; + 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; @@ -678,22 +671,14 @@ void fused_attn_arbitrary_seqlen_bwd_impl( bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); - bool is_bottom_right = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - if (is_bottom_right && s_q == s_kv && !is_padding) { - is_causal = true; - is_bottom_right = false; - bottom_right_diagonal = false; - } + bool is_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (dropout_probability != 0.0f); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + NVTE_QKV_Format q_format = cfg.q_format; + NVTE_QKV_Format kv_format = cfg.kv_format; + bool is_ragged_q = cfg.is_ragged_q; + bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); @@ -752,7 +737,16 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - if (it != cache.end()) { + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] + sm_arch_ != 120) { // [GRAPH-DEBUG] + // The backward impl has no cu_seqlens-directly path; it always buckets the batch. + fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] + "bwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), + /*legacy=*/true); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } @@ -879,7 +873,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); } - if (is_causal || is_bottom_right) { + if (is_causal || is_causal_bottom_right) { sdpa_backward_options.set_diagonal_band_right_bound(0); } @@ -1365,6 +1359,7 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = true; graph_cfg.derive(); size_t workspace_size = 0; @@ -1389,6 +1384,7 @@ std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handl std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = false; graph_cfg.derive(); size_t workspace_size = 0; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index db919c92d6..1afa22b73e 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -41,7 +41,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de const int64_t d_qk = static_cast(cfg.head_dim_qk); const int64_t d_v = static_cast(cfg.head_dim_v); const bool is_training = cfg.is_training; - const float scaling_factor = cfg.attn_scale; + 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; @@ -58,8 +58,8 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + bool is_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; bool is_dropout = (is_training && dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; @@ -134,7 +134,9 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - if (it != cache.end()) { + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } @@ -197,10 +199,10 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de } else if (is_mxfp8) { NVTE_QKV_Format q_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format - : nvte_get_q_format(qkv_layout); + : cfg.q_format; NVTE_QKV_Format kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format - : nvte_get_kv_format(qkv_layout); + : cfg.kv_format; std::vector q_scale_strides(4); std::vector k_scale_strides(4); std::vector v_scale_strides(4); @@ -238,6 +240,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de sdpa_options = fe::graph::SDPA_fp8_attributes() .set_name("sdpa_fp8") .set_generate_stats(true) + .set_causal_mask(is_causal) .set_attn_scale(attn_scale); fe::DiagonalAlignment_t const& diagonal_alignment = @@ -253,7 +256,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de sdpa_options.set_diagonal_band_right_bound(window_size_right); } } - if (is_causal) { + if (is_causal_bottom_right) { sdpa_options.set_diagonal_band_right_bound(0); } @@ -512,7 +515,7 @@ void fused_attn_fp8_bwd_impl( const int64_t s_kv = static_cast(cfg.max_seqlen_kv); const int64_t d_qk = static_cast(cfg.head_dim_qk); const int64_t d_v = static_cast(cfg.head_dim_v); - const float scaling_factor = cfg.attn_scale; + 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; @@ -533,8 +536,8 @@ void fused_attn_fp8_bwd_impl( bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + bool is_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; bool is_dropout = (dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; @@ -615,7 +618,9 @@ void fused_attn_fp8_bwd_impl( auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - if (it != cache.end()) { + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } @@ -712,8 +717,8 @@ void fused_attn_fp8_bwd_impl( scale_dV = mha_graph->tensor(1.0f); } } else if (is_mxfp8) { - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + NVTE_QKV_Format q_format = cfg.q_format; + NVTE_QKV_Format kv_format = cfg.kv_format; NVTE_QKV_Format q_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : q_format; NVTE_QKV_Format kv_scale_inv_format = @@ -817,6 +822,7 @@ void fused_attn_fp8_bwd_impl( fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; sdpa_backward_options = fe::graph::SDPA_fp8_backward_attributes() .set_name("sdpa_fp8_backward") + .set_causal_mask(is_causal) .set_attn_scale(attn_scale); fe::DiagonalAlignment_t const& diagonal_alignment = @@ -832,7 +838,7 @@ void fused_attn_fp8_bwd_impl( sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); } } - if (is_causal) { + if (is_causal_bottom_right) { sdpa_backward_options.set_diagonal_band_right_bound(0); } @@ -1347,6 +1353,7 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = true; graph_cfg.derive(); size_t workspace_size = 0; @@ -1372,6 +1379,7 @@ std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handl std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = false; graph_cfg.derive(); size_t workspace_size = 0; diff --git a/transformer_engine/common/fused_attn/graph_debug.h b/transformer_engine/common/fused_attn/graph_debug.h index c5afc8f035..c1a60a39af 100644 --- a/transformer_engine/common/fused_attn/graph_debug.h +++ b/transformer_engine/common/fused_attn/graph_debug.h @@ -9,11 +9,22 @@ // // Counts fused-attention cuDNN graph *builds* (cache misses that construct a new // graph) vs. *executions* (real forward/backward runs, excluding workspace-sizing -// probes) to detect redundant graph construction. +// probes) to detect redundant graph construction. Also logs every graph-cache +// lookup (HIT/MISS + the key fields) to diagnose stale-cache reuse across tests. // // Enable at runtime with: export NVTE_FUSED_ATTN_GRAPH_DEBUG=1 -// A running "BUILD" line is printed whenever a new graph is constructed, and a -// "SUMMARY" line with final totals is printed at process exit. +// - A "BUILD" line is printed whenever a new graph is constructed. +// - A "HIT"/"MISS" line with the key fields is printed on every cache lookup. +// - A "thd ... path=legacy|direct" line is printed on every THD (ragged) lookup, showing +// which impl path (bucketed batch vs. real batch) the graph was built for. +// - A "SUMMARY" line with final build/exec totals is printed at process exit, followed by a +// "THD-PATH" line with per-path lookup/build totals (low builds/lookups on the legacy path +// means batch bucketing is collapsing distinct batch sizes onto shared graphs). +// +// Separately, force every lookup to miss (never reuse a cached graph) with: +// export NVTE_FUSED_ATTN_DISABLE_CACHE=1 +// If a suite that fails with the cache enabled passes with it disabled, the bug +// is stale-cache reuse (an incomplete make_cache_key / operator<). // // To remove all of this instrumentation later: // 1. Delete this file (graph_debug.h). @@ -29,10 +40,22 @@ #include #include #include +#include + +#include "config_and_params.h" // [GRAPH-DEBUG] for FusedAttnConfig field dump namespace transformer_engine { namespace fused_attn_graph_debug { +// Short, stable per-thread id (0, 1, 2, ...) assigned on first use. The graph caches are +// static thread_local, so a graph built on one thread is invisible to another; tagging every +// lookup with its thread id makes cross-thread rebuilds of an identical key visible. +inline unsigned thread_seq_id() { + static std::atomic next{0}; + static thread_local unsigned id = next.fetch_add(1); + return id; +} + inline std::atomic &fwd_built() { static std::atomic v{0}; return v; @@ -50,6 +73,29 @@ inline std::atomic &bwd_exec() { return v; } +// THD (ragged) cache lookups split by which impl path the graph was built for: +// legacy = batch quantized into a bucket (many batch sizes share one graph) +// direct = cu_seqlens fed to cuDNN directly (real batch baked in, no batch sharing) +// "builds" counts the lookups that actually constructed a new graph. A low builds/lookups +// ratio on the legacy path is the visible sign that batch bucketing is collapsing distinct +// batch sizes onto shared graphs. +inline std::atomic &thd_legacy_lookup() { + static std::atomic v{0}; + return v; +} +inline std::atomic &thd_legacy_build() { + static std::atomic v{0}; + return v; +} +inline std::atomic &thd_direct_lookup() { + static std::atomic v{0}; + return v; +} +inline std::atomic &thd_direct_build() { + static std::atomic v{0}; + return v; +} + inline bool enabled() { static const bool on = [] { const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG"); @@ -60,18 +106,32 @@ inline bool enabled() { inline void dump(const char *event) { std::fprintf(stderr, - "[GRAPH-DEBUG] %-10s | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", event, - static_cast(fwd_built().load()), + "[GRAPH-DEBUG] %-10s | tid=%u | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", + event, thread_seq_id(), static_cast(fwd_built().load()), static_cast(fwd_exec().load()), static_cast(bwd_built().load()), static_cast(bwd_exec().load())); std::fflush(stderr); } +inline void dump_thd_summary() { + std::fprintf( + stderr, + "[GRAPH-DEBUG] THD-PATH | legacy lookups=%llu builds=%llu | direct lookups=%llu builds=%llu\n", + static_cast(thd_legacy_lookup().load()), + static_cast(thd_legacy_build().load()), + static_cast(thd_direct_lookup().load()), + static_cast(thd_direct_build().load())); + std::fflush(stderr); +} + inline void register_summary_once() { static const bool registered = [] { std::atexit([] { - if (enabled()) dump("SUMMARY"); + if (enabled()) { + dump("SUMMARY"); + dump_thd_summary(); + } }); return true; }(); @@ -101,6 +161,83 @@ inline void note_bwd_exec() { bwd_exec().fetch_add(1); } +// Returns true when the graph cache should be bypassed (every lookup treated as a +// miss so a fresh graph is built each call). Gated by NVTE_FUSED_ATTN_DISABLE_CACHE. +inline bool cache_disabled() { + static const bool off = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_DISABLE_CACHE"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return off; +} + +// Logs one graph-cache lookup with its outcome (HIT/MISS) and the *real* (pre- +// normalization) config fields. A std::map HIT means the two configs compare equal +// under operator<, so the field that actually distinguishes a wrongly-reused graph +// is one that make_cache_key() normalized away or that operator< omits -- pass the +// real cfg (not the normalized cache key) here so that difference is visible when +// diffing a wrong HIT against the earlier BUILD that created the reused graph. +inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { + if (!enabled()) return; + register_summary_once(); + std::fprintf( + stderr, + "[GRAPH-DEBUG] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld bias=%lld " + "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " + "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " + "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " + "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " + "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", + pass, hit ? "HIT" : "MISS", + (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", thread_seq_id(), + static_cast(c.is_training), + static_cast(c.deterministic), static_cast(c.cuda_graph), + static_cast(c.return_max_logit), static_cast(c.is_forward), + static_cast(c.attn_mask_type), static_cast(c.bias_type), + static_cast(c.window_size_left), static_cast(c.window_size_right), + static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), + static_cast(c.scaling_mode), static_cast(c.dropout), + static_cast(c.attn_scale), static_cast(c.qkv_dtype), + static_cast(c.o_dtype), static_cast(c.do_dtype), + static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), + static_cast(c.o_format), static_cast(c.do_format), + static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), + static_cast(c.do_scale_inv_format), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.head_dim_qk), static_cast(c.head_dim_v), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), + static_cast(c.num_pages_v), static_cast(c.page_size_k), + static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), + static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), + static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), + static_cast(c.bias_seqlen_kv)); + std::fflush(stderr); +} + +// Records, for one THD (ragged) cache lookup, which impl path the graph was built for -- +// "legacy" (batch quantized into a bucket) vs "direct" (real batch fed via cu_seqlens) -- and +// whether it hit the cache. `built` should reflect whether a new graph was actually constructed +// (i.e. a real miss, or a hit forced to rebuild by NVTE_FUSED_ATTN_DISABLE_CACHE). Comparing +// per-path lookups vs builds in the THD-PATH summary shows the batch-bucketing effect. +inline void note_thd_lookup(const char *pass, bool hit, bool built, bool legacy) { + if (!enabled()) return; + register_summary_once(); + if (legacy) { + thd_legacy_lookup().fetch_add(1); + if (built) thd_legacy_build().fetch_add(1); + } else { + thd_direct_lookup().fetch_add(1); + if (built) thd_direct_build().fetch_add(1); + } + std::fprintf(stderr, "[GRAPH-DEBUG] thd %-3s %-4s | tid=%u | path=%s%s\n", pass, + hit ? "HIT" : "MISS", thread_seq_id(), legacy ? "legacy" : "direct", + (hit && built) ? " [cache-disabled->rebuild]" : ""); + std::fflush(stderr); +} + } // namespace fused_attn_graph_debug } // namespace transformer_engine From ade19fe230d0e312b533cf15003b41aa2b9e6db7 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Sat, 18 Jul 2026 08:43:12 -0700 Subject: [PATCH 32/63] match fused attn availability probe to runtime for FP8 specs and per-step CP configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 4 + .../attention/test_attention_with_cp.py | 4 + tests/pytorch/utils.py | 6 + .../dot_product_attention/context_parallel.py | 77 +++++++ .../dot_product_attention.py | 4 + .../attention/dot_product_attention/utils.py | 210 ++++++++++++++---- 6 files changed, 260 insertions(+), 45 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 97a4e97893..e63b7b7b04 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1784,6 +1784,7 @@ def test_dpa_fp8_extra_state(model, dtype): available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout="sb3hd", is_training=is_training, deterministic=_deterministic, @@ -2014,6 +2015,7 @@ def test_mha_fp8_vs_f16( available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout=qkv_format.replace("hd", "h3d"), fp8=True, fp8_meta=fp8_meta, @@ -2271,6 +2273,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout=qkv_layout, fp8=True, fp8_meta=fp8_meta, @@ -2593,6 +2596,7 @@ def test_custom_mha_fp8_vs_f16(dtype, model): available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout="bs3hd", fp8=True, fp8_meta=fp8_meta, diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 681ee5e6e0..9ffdb865fa 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -378,6 +378,8 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type config, qkv_dtype=dtypes[dtype], qkv_layout="_".join([qkv_format] * 3), + cp_size=num_gpus, + cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, ) flash_attn_supported, *_ = available_backends if not flash_attn_supported: @@ -638,6 +640,8 @@ def test_cp_with_fused_attention( fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, + cp_size=num_gpus, + cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, ) _, fused_attn_supported, _ = available_backends diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 4aed95cb2c..34bfe7b939 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -339,6 +339,7 @@ def get_available_attention_backends( config: ModelConfig, qkv_dtype: torch.dtype, qkv_layout: str, + nominal_dtype: Optional[torch.dtype] = None, pad_between_seqs: bool = False, deterministic: bool = False, fp8: bool = False, @@ -347,6 +348,8 @@ def get_available_attention_backends( inference_params: Optional[InferenceParams] = None, score_mod: bool = False, score_mod_bprop: bool = False, + cp_size: int = 1, + cp_size_a2a: int = 1, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration""" @@ -389,6 +392,7 @@ def get_available_attention_backends( def test(): attention_params = AttentionParams( qkv_dtype=qkv_dtype, + nominal_dtype=nominal_dtype, qkv_layout=qkv_layout, batch_size=config.batch_size, num_heads=config.num_heads, @@ -408,6 +412,8 @@ def test(): attention_dropout=config.dropout_p, context_parallel=config.context_parallel, cp_comm_type=config.cp_comm_type, + cp_size=cp_size, + cp_size_a2a=cp_size_a2a, deterministic=deterministic, fp8=fp8, fp8_meta=fp8_meta, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index a62ca73187..5c8f01366d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4918,6 +4918,83 @@ def backward(ctx, dout, *_args): ) +def cp_per_step_configs( + cp_comm_type, + cp_size, + cp_size_a2a, + *, + max_seqlen_q, + max_seqlen_kv, + num_heads, + num_gqa_groups, + attn_mask_type, + window_size, + bottom_right_diagonal, +): + """Per-step attention configs a context-parallel run dispatches to its attention backend. + + CP runs attention in multiple steps, each with a distinct config (e.g. mask, and seqlens) + that differs from the single global config. This function returns the list of those distinct + per-step configs so `get_attention_backend` can check if the backend supports all of them. + """ + is_causal = "causal" in attn_mask_type + padding_or_no_mask = "padding" if "padding" in attn_mask_type else "no_mask" + window_left, window_right = window_size + + def config(mask, s_q, s_kv, heads, gqa, bottom_right): + return dict( + attn_mask_type=mask, + max_seqlen_q=s_q, + max_seqlen_kv=s_kv, + num_attn_heads=heads, + num_gqa_groups=gqa, + window_size_left=window_left, + window_size_right=window_right, + bottom_right_diagonal=bottom_right, + ) + + if cp_comm_type == "a2a": + # split heads across the cp ranks + return [ + config( + attn_mask_type, + max_seqlen_q, + max_seqlen_kv, + num_heads // cp_size, + num_gqa_groups // cp_size, + bottom_right_diagonal, + ) + ] + + if cp_comm_type == "all_gather": + # one short Q chunk vs a growing KV chunk; causal -> causal_bottom_right + s_q = max_seqlen_q // (2 * cp_size) + s_kv_chunk = max_seqlen_kv // (2 * cp_size) + mask, br = attn_mask_type, bottom_right_diagonal + if is_causal and "bottom_right" not in attn_mask_type: + mask, br = attn_mask_type + "_bottom_right", True + # s_kv ranges from s_kv_chunk, i*s_kv_chunk, ..., max_seqlen_kv + # check a single chunk and the full KV + return [ + config(mask, s_q, s_kv, num_heads, num_gqa_groups, br) + for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) + ] + + # p2p and a2a+p2p: split heads across the a2a subgroup, and ring over the p2p subgroup + p2p_size = cp_size // cp_size_a2a + heads = num_heads // cp_size_a2a + gqa = num_gqa_groups // cp_size_a2a + r_q = max_seqlen_q // p2p_size + r_kv = max_seqlen_kv // p2p_size + if not is_causal: + return [config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal)] + return [ + config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal), # diagonal + config(padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal), # lower-triangle + config(padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal), # upper-triangle + ] + + def attn_forward_func_with_cp( is_training, q, 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 c293aeae88..8c2e181eb2 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 @@ -1449,11 +1449,14 @@ def forward( # adjust max_seqlen and cu_seqlens for CP cp_size = 1 + cp_size_a2a = 1 if isinstance(self.cp_group, dist_group_type): cp_size = get_distributed_world_size(self.cp_group) elif isinstance(self.cp_group, list): for group in self.cp_group: cp_size *= get_distributed_world_size(group) + if self.cp_comm_type == "a2a+p2p" and len(self.cp_group) > 0: + cp_size_a2a = get_distributed_world_size(self.cp_group[0]) context_parallel = cp_size > 1 if q_format in ["sbhd", "bshd"]: max_seqlen_q *= cp_size @@ -1608,6 +1611,7 @@ def forward( context_parallel=context_parallel, cp_comm_type=self.cp_comm_type, cp_size=cp_size, + cp_size_a2a=cp_size_a2a, deterministic=self.deterministic, is_training=self.training, fp8=self.fp8, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 398f61bee9..dcb4db1e5e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -199,6 +199,9 @@ class AttentionParams: Type of query/key/value tensors, {`torch.Tensor`, `Float8Tensor`}. qkv_dtype : torch.dtype, default = torch.bfloat16 Data type of query/key/value tensors. + nominal_dtype : Optional[torch.dtype], default = None + Model precision (F16/BF16) of the unquantized tensors (O, and dQ/dK/dV under + current/mxfp8) when `qkv_dtype` itself is FP8. qkv_layout : str, default = "sbh3d" Query/key/value tensor memory layout. batch_size : int, default = 1 @@ -245,7 +248,9 @@ class AttentionParams: cp_comm_type : str, default = "p2p" The communication type of context parallelism. cp_size : int, default = 1 - The group size of context parallelism. + The (total) group size of context parallelism. + cp_size_a2a : int, default = 1 + The all-to-all subgroup size when `cp_comm_type == "a2a+p2p"`. deterministic : bool, default = False Whether to run `DotProductAttention` with determinism or not. is_training : bool, default = True @@ -278,6 +283,7 @@ class AttentionParams: qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor qkv_dtype: torch.dtype = torch.bfloat16 + nominal_dtype: Optional[torch.dtype] = None qkv_layout: str = "sbh3d" batch_size: int = 1 num_heads: int = 16 @@ -300,6 +306,7 @@ class AttentionParams: context_parallel: bool = False cp_comm_type: str = "p2p" cp_size: int = 1 + cp_size_a2a: int = 1 deterministic: bool = False is_training: bool = True fp8: bool = False @@ -424,6 +431,7 @@ def get_attention_backend( # is shifted over to the caller of this function qkv_type = attention_params.qkv_type qkv_dtype = attention_params.qkv_dtype + nominal_dtype = attention_params.nominal_dtype qkv_layout = attention_params.qkv_layout batch_size = attention_params.batch_size num_heads = attention_params.num_heads @@ -445,7 +453,8 @@ def get_attention_backend( attention_dropout = attention_params.attention_dropout context_parallel = attention_params.context_parallel cp_comm_type = attention_params.cp_comm_type - cp_size = attention_params.cp_size # pylint: disable=unused-variable + cp_size = attention_params.cp_size + cp_size_a2a = attention_params.cp_size_a2a deterministic = attention_params.deterministic is_training = attention_params.is_training fp8 = attention_params.fp8 @@ -1448,39 +1457,18 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # Filter: cuDNN support fused_attention_backend = None if use_fused_attention: - # ``DType`` is implicitly convertible to ``transformer_engine::DType`` - # on the C++ side, so pass it straight to the pybind function. - qkv_type = TE_DType[qkv_dtype] - o_type = qkv_type - do_type = qkv_type - dqkv_type = qkv_type - scaling_mode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING - qkv_scale_inv_format = None - do_scale_inv_format = None - if fp8 and fp8_meta["recipe"].fp8_dpa: - recipe = fp8_meta["recipe"] - qkv_type = get_fp8_te_dtype(recipe, fprop_tensor=True) - 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 = qkv_type - do_type = o_type - dqkv_type = qkv_type - o_format = q_format - do_format = o_format - dqkv_layout = qkv_layout + recipe = fp8_meta["recipe"] if (fp8 and fp8_meta["recipe"].fp8_dpa) else None + cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" + spec = get_fused_attn_spec( + recipe, qkv_dtype, qkv_layout, cs_o_in_f16=cs_o_in_f16, nominal_dtype=nominal_dtype + ) + qkv_type, o_type, do_type, dqkv_type = spec.qkv, spec.o, spec.do, spec.dqkv + scaling_mode = spec.scaling_mode + qkv_scale_inv_format = spec.scale_inv_format + do_scale_inv_format = spec.scale_inv_format + o_format = spec.o_format + do_format = spec.do_format + dqkv_layout = spec.dqkv_layout num_pages_k = num_pages_v = 0 page_size_k = page_size_v = 0 max_pages_per_seq_k = max_pages_per_seq_v = 0 @@ -1491,7 +1479,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt bias_batch_size = bias_num_heads = bias_seqlen_q = bias_seqlen_kv = 0 if fu_core_attention_bias_shape is not None: bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = fu_core_attention_bias_shape - fused_attn_params = FusedAttentionParams( + base_fused_attn_kwargs = dict( is_training=is_training, deterministic=deterministic, cuda_graph=cuda_graph, @@ -1509,7 +1497,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt o_dtype=o_type, do_dtype=do_type, dqkv_dtype=dqkv_type, - qkv_layout=QKVLayout[qkv_layout], + qkv_layout=QKVLayout[spec.qkv_layout], o_format=QKVFormat[o_format], do_format=QKVFormat[do_format], dqkv_layout=QKVLayout[dqkv_layout], @@ -1535,15 +1523,65 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, ) - fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: - logger.debug( - "Disabling FusedAttention: %s", - reject_message, + + if context_parallel: + from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( + cp_per_step_configs, ) - use_fused_attention = False - fused_attention_backend = None - elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]: + + per_step_configs = cp_per_step_configs( + cp_comm_type, + cp_size, + cp_size_a2a, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + num_heads=num_heads, + num_gqa_groups=num_gqa_groups, + attn_mask_type=attn_mask_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + ) + else: + per_step_configs = [None] + + for step_config in per_step_configs: + fused_attn_kwargs = dict(base_fused_attn_kwargs) + if step_config is not None: + step_seqlen_q = step_config["max_seqlen_q"] + step_seqlen_kv = step_config["max_seqlen_kv"] + fused_attn_kwargs.update( + attn_mask_type=AttnMaskType[step_config["attn_mask_type"]], + max_seqlen_q=step_seqlen_q, + max_seqlen_kv=step_seqlen_kv, + num_attn_heads=step_config["num_attn_heads"], + num_gqa_groups=step_config["num_gqa_groups"], + window_size_left=step_config["window_size_left"], + window_size_right=step_config["window_size_right"], + bottom_right_diagonal=step_config["bottom_right_diagonal"], + ) + if bias_seqlen_q != 1: + fused_attn_kwargs["bias_seqlen_q"] = step_seqlen_q + if bias_seqlen_kv != 1: + fused_attn_kwargs["bias_seqlen_kv"] = step_seqlen_kv + fused_attn_params = FusedAttentionParams(**fused_attn_kwargs) + fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) + if fused_attention_backend == FusedAttnBackend["No_Backend"]: + logger.debug( + "Disabling FusedAttention: %s%s", + reject_message, + f" (context-parallel per-step config {step_config})" + if step_config is not None + else "", + ) + use_fused_attention = False + fused_attention_backend = None + break + + if ( + use_fused_attention + and has_score_mod + and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"] + ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " "F16/BF16 arbitrary-seqlen", @@ -2403,6 +2441,88 @@ def get_qkv_format( return qkv_format, q_format, kv_format +@dataclass(frozen=True) +class FusedAttnSpec: + """Fused-attention spec for a given config. + + Mirrors what `FusedAttnFunc` feeds `fused_attn_fwd`/`fused_attn_bwd` (backends.py), + so the availability probe (`get_attention_backend`) cannot drift from runtime. + """ + + scaling_mode: Any + qkv: Any + o: Any + do: Any + dqkv: Any + scale_inv_format: Optional[str] + qkv_layout: str + o_format: str + do_format: str + dqkv_layout: str + + +def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_dtype=None): + """Resolve fused-attention specs, e.g. tensor dtypes, formats, for a given config. + + `nominal_dtype` is the model precision (F16/BF16) of the tensors that stay unquantized in + FP8 attention (O, and dQ/dK/dV under current/mxfp8). It is only consulted when `qkv_dtype` itself is FP8. + """ + q_format = get_qkv_format(qkv_layout)[1] + eff_qkv_layout = qkv_layout # FP16/BF16 + if recipe is not None: + if not recipe.mxfp8(): + # Delayed/current scaling + eff_qkv_layout = qkv_layout.replace("paged_kv_", "") + elif qkv_layout in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): + eff_qkv_layout = qkv_layout # MXFP8 fast path + else: + eff_qkv_layout = "bhsd_bhsd_bhsd" # MXFP8 slow path + layout_kwargs = dict( + qkv_layout=eff_qkv_layout, + o_format=q_format, + do_format=q_format, + dqkv_layout=qkv_layout, + ) + + if qkv_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + ref = TE_DType[nominal_dtype if nominal_dtype is not None else torch.bfloat16] + else: + ref = TE_DType[qkv_dtype] + + # FP16/BF16: every tensor is in model precision; scaling_mode is a placeholder + if recipe is None: + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, ref, ref, ref, ref, None, + **layout_kwargs, + ) + + fprop_fp8 = get_fp8_te_dtype(recipe, fprop_tensor=True) + grad_fp8 = get_fp8_te_dtype(recipe, fprop_tensor=False) + + # MXFP8 block scaling: Q/K/V/dO are in MXFP8; O/dQ/dK/dV stay in model precision + if recipe.mxfp8(): + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, fprop_fp8, ref, grad_fp8, ref, "bhsd", + **layout_kwargs, + ) + + # FP8 current scaling: Q/K/V/dO are in FP8; O in model precision if `cs_o_in_f16` (default), otherwise FP8; + # dQ/dK/dV in model precision + if recipe.float8_current_scaling(): + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, + fprop_fp8, ref if cs_o_in_f16 else fprop_fp8, grad_fp8, ref, None, + **layout_kwargs, + ) + + # FP8 delayed scaling: Q/K/V/O are in FP8 (e.g. E4M3); dO/dQ/dK/dV in FP8 (e.g. E5M2) + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, + fprop_fp8, fprop_fp8, grad_fp8, grad_fp8, None, + **layout_kwargs, + ) + + def get_qkv_layout( q: torch.Tensor, k: torch.Tensor, From 42e97474d6a2042d8519134d10b6caf7a58c2645 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Sat, 18 Jul 2026 08:45:32 -0700 Subject: [PATCH 33/63] fix bias for jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn.py | 16 ++++++++++++++++ .../common/fused_attn/config_and_params.cpp | 2 +- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 1768e0227d..b53ce95668 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -520,6 +520,18 @@ def _check_configs(self): "is either BSHD_BSHD_BSHD or THD_THD_THD" ) + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None + if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + if self.bias_shape == BiasShape._1HSS: + bias_batch, bias_heads = 1, self.num_heads_q + elif self.bias_shape == BiasShape._B1SS: + bias_batch, bias_heads = self.batch_size, 1 + elif self.bias_shape == BiasShape._BHSS: + bias_batch, bias_heads = self.batch_size, self.num_heads_q + elif self.bias_shape == BiasShape._11SS: + bias_batch, bias_heads = 1, 1 + bias_seqlen_q, bias_seqlen_kv = self.max_seqlen_q, self.max_seqlen_kv + self.backend, message = FusedAttnHelper( self.is_training, self.batch_size, @@ -538,6 +550,10 @@ def _check_configs(self): self.head_dim_v, (-1, -1) if self.window_size is None else self.window_size, self.attn_mask_type.is_bottom_right(), + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ).get_fused_attn_backend() if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: pytest.skip(message) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 40b7831ace..ccb9e6fc3b 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -245,7 +245,7 @@ FusedAttnConfig FusedAttnFwdParams::make_config() const { cfg.num_tokens_kv = t_kv; if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && - input_Bias->data.dptr != nullptr && input_Bias->data.shape.size() >= 4) { + input_Bias->data.shape.size() >= 4) { cfg.bias_batch_size = input_Bias->data.shape[0]; cfg.bias_num_heads = input_Bias->data.shape[1]; cfg.bias_seqlen_q = input_Bias->data.shape[2]; From c66028a8dec16832a23697fb9653296aa34980a4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 15:45:55 +0000 Subject: [PATCH 34/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/config_and_params.cpp | 9 +-- .../common/fused_attn/config_and_params.h | 21 +++--- .../fused_attn_f16_arbitrary_seqlen.cu | 32 ++++----- .../common/fused_attn/fused_attn_fp8.cu | 13 ++-- .../common/fused_attn/graph_debug.h | 69 ++++++++++--------- .../dot_product_attention/context_parallel.py | 8 ++- .../attention/dot_product_attention/utils.py | 51 ++++++++++---- 7 files changed, 117 insertions(+), 86 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index ccb9e6fc3b..16727f80e8 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -52,8 +52,9 @@ void FusedAttnConfig::derive() { (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); is_causal = (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK); - is_causal_bottom_right = (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + is_causal_bottom_right = + (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); // bucket the THD (ragged) batch and token counts const size_t tokens_q = num_tokens_q != 0 ? num_tokens_q : static_cast(b * sq); @@ -104,8 +105,8 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { const bool has_window = cache_cfg.window_size_left != -1 || cache_cfg.window_size_right != -1; if (!cache_cfg.is_causal && !cache_cfg.is_causal_bottom_right && !has_window) { cache_cfg.bottom_right_diagonal = false; - } else if (cache_cfg.is_causal_bottom_right && cache_cfg.max_seqlen_q == cache_cfg.max_seqlen_kv && - !cache_cfg.is_padding) { + } else if (cache_cfg.is_causal_bottom_right && + cache_cfg.max_seqlen_q == cache_cfg.max_seqlen_kv && !cache_cfg.is_padding) { // square bottom-right causal collapses to top-left causal (mirrors the impl). cache_cfg.bottom_right_diagonal = false; } diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 5ec6b02822..4839a73af7 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -151,23 +151,22 @@ struct FusedAttnConfig { bias_type, window_size_left, window_size_right, bottom_right_diagonal, softmax_type, scaling_mode, dropout, attn_scale, qkv_dtype, o_dtype, do_dtype, dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, - head_dim_qk, head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, - num_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_batch_size, bias_num_heads, - bias_seqlen_q, bias_seqlen_kv) < + do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, + head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_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_batch_size, bias_num_heads, bias_seqlen_q, + bias_seqlen_kv) < std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, rhs.attn_scale, rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, - rhs.num_attn_heads, - rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, rhs.max_seqlen_q, - rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, 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_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, + rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, + 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); } // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields 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 09fae572fb..bb593512f5 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 @@ -188,15 +188,15 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] - fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] - sm_arch_ != 120) { // [GRAPH-DEBUG] - fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] + sm_arch_ != 120) { // [GRAPH-DEBUG] + fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] "fwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), - /*legacy=*/!use_cu_seqlens_directly); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + /*legacy=*/!use_cu_seqlens_directly); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } @@ -737,16 +737,16 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] - fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] - sm_arch_ != 120) { // [GRAPH-DEBUG] + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] + sm_arch_ != 120) { // [GRAPH-DEBUG] // The backward impl has no cu_seqlens-directly path; it always buckets the batch. - fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] + fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] "bwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), - /*legacy=*/true); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + /*legacy=*/true); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 1afa22b73e..5163a3601a 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -134,9 +134,9 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } @@ -197,9 +197,8 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de scale_o = mha_graph->tensor(1.0f); } } else if (is_mxfp8) { - NVTE_QKV_Format q_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) - ? qkv_scale_inv_format - : cfg.q_format; + NVTE_QKV_Format q_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.q_format; NVTE_QKV_Format kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.kv_format; @@ -618,9 +617,9 @@ void fused_attn_fp8_bwd_impl( auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { // if hit, return auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] auto graph = it->second; return graph; } diff --git a/transformer_engine/common/fused_attn/graph_debug.h b/transformer_engine/common/fused_attn/graph_debug.h index c1a60a39af..8a69b85e78 100644 --- a/transformer_engine/common/fused_attn/graph_debug.h +++ b/transformer_engine/common/fused_attn/graph_debug.h @@ -105,23 +105,24 @@ inline bool enabled() { } inline void dump(const char *event) { - std::fprintf(stderr, - "[GRAPH-DEBUG] %-10s | tid=%u | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", - event, thread_seq_id(), static_cast(fwd_built().load()), - static_cast(fwd_exec().load()), - static_cast(bwd_built().load()), - static_cast(bwd_exec().load())); + std::fprintf( + stderr, + "[GRAPH-DEBUG] %-10s | tid=%u | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", event, + thread_seq_id(), static_cast(fwd_built().load()), + static_cast(fwd_exec().load()), + static_cast(bwd_built().load()), + static_cast(bwd_exec().load())); std::fflush(stderr); } inline void dump_thd_summary() { - std::fprintf( - stderr, - "[GRAPH-DEBUG] THD-PATH | legacy lookups=%llu builds=%llu | direct lookups=%llu builds=%llu\n", - static_cast(thd_legacy_lookup().load()), - static_cast(thd_legacy_build().load()), - static_cast(thd_direct_lookup().load()), - static_cast(thd_direct_build().load())); + std::fprintf(stderr, + "[GRAPH-DEBUG] THD-PATH | legacy lookups=%llu builds=%llu | direct lookups=%llu " + "builds=%llu\n", + static_cast(thd_legacy_lookup().load()), + static_cast(thd_legacy_build().load()), + static_cast(thd_direct_lookup().load()), + static_cast(thd_direct_build().load())); std::fflush(stderr); } @@ -182,32 +183,32 @@ inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig register_summary_once(); std::fprintf( stderr, - "[GRAPH-DEBUG] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld bias=%lld " + "[GRAPH-DEBUG] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld " + "bias=%lld " "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", - pass, hit ? "HIT" : "MISS", - (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", thread_seq_id(), - static_cast(c.is_training), - static_cast(c.deterministic), static_cast(c.cuda_graph), - static_cast(c.return_max_logit), static_cast(c.is_forward), - static_cast(c.attn_mask_type), static_cast(c.bias_type), - static_cast(c.window_size_left), static_cast(c.window_size_right), - static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), - static_cast(c.scaling_mode), static_cast(c.dropout), - static_cast(c.attn_scale), static_cast(c.qkv_dtype), - static_cast(c.o_dtype), static_cast(c.do_dtype), - static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), - static_cast(c.o_format), static_cast(c.do_format), - static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), - static_cast(c.do_scale_inv_format), static_cast(c.batch_size), - static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), - static_cast(c.head_dim_qk), static_cast(c.head_dim_v), - static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), - static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), - static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + pass, hit ? "HIT" : "MISS", (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", + thread_seq_id(), static_cast(c.is_training), static_cast(c.deterministic), + static_cast(c.cuda_graph), static_cast(c.return_max_logit), + static_cast(c.is_forward), static_cast(c.attn_mask_type), + static_cast(c.bias_type), static_cast(c.window_size_left), + static_cast(c.window_size_right), static_cast(c.bottom_right_diagonal), + static_cast(c.softmax_type), static_cast(c.scaling_mode), + static_cast(c.dropout), static_cast(c.attn_scale), + static_cast(c.qkv_dtype), static_cast(c.o_dtype), + static_cast(c.do_dtype), static_cast(c.dqkv_dtype), + static_cast(c.qkv_layout), static_cast(c.o_format), + static_cast(c.do_format), static_cast(c.dqkv_layout), + static_cast(c.qkv_scale_inv_format), static_cast(c.do_scale_inv_format), + static_cast(c.batch_size), static_cast(c.num_attn_heads), + static_cast(c.num_gqa_groups), static_cast(c.head_dim_qk), + static_cast(c.head_dim_v), static_cast(c.max_seqlen_q), + static_cast(c.max_seqlen_kv), static_cast(c.num_tokens_q), + static_cast(c.num_tokens_kv), static_cast(c.bucketed_batch_size), + static_cast(c.bucketed_num_tokens_q), static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), static_cast(c.num_pages_v), static_cast(c.page_size_k), static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 5c8f01366d..8b8bb1b779 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4990,8 +4990,12 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right): return [config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal)] return [ config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal), # diagonal - config(padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal), # lower-triangle - config(padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal), # upper-triangle + config( + padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal + ), # lower-triangle + config( + padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal + ), # upper-triangle ] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index dcb4db1e5e..c3b9c1da1b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1426,7 +1426,10 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fu_core_attention_bias_shape = (batch_size, num_heads, max_seqlen_q, max_seqlen_kv) fu_core_attention_bias_shape_type = None - if fu_core_attention_bias_type == "post_scale_bias" and fu_core_attention_bias_shape is not None: + if ( + fu_core_attention_bias_type == "post_scale_bias" + and fu_core_attention_bias_shape is not None + ): b, h, sq, _skv = fu_core_attention_bias_shape if b == batch_size and h == num_heads: fu_core_attention_bias_shape_type = "bhss" @@ -1439,7 +1442,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt else: raise ValueError( "core_attention_bias tensor must be in one of " - f'{{"bhss", "1hss", "b1ss", "11ss", "111s"}} shapes. ' + '{"bhss", "1hss", "b1ss", "11ss", "111s"} shapes. ' f"Found (b,h,sq,skv) = ({b},{h},{sq},{_skv})" ) if ( @@ -1475,10 +1478,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if inference_params is not None and getattr(inference_params, "is_paged", False): num_pages_k = num_pages_v = inference_params.total_num_pages page_size_k = page_size_v = inference_params.page_size - max_pages_per_seq_k = max_pages_per_seq_v = inference_params.cache_manager.max_pages_per_seq + max_pages_per_seq_k = max_pages_per_seq_v = ( + inference_params.cache_manager.max_pages_per_seq + ) bias_batch_size = bias_num_heads = bias_seqlen_q = bias_seqlen_kv = 0 if fu_core_attention_bias_shape is not None: - bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = fu_core_attention_bias_shape + bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = ( + fu_core_attention_bias_shape + ) base_fused_attn_kwargs = dict( is_training=is_training, deterministic=deterministic, @@ -1569,9 +1576,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug( "Disabling FusedAttention: %s%s", reject_message, - f" (context-parallel per-step config {step_config})" - if step_config is not None - else "", + ( + f" (context-parallel per-step config {step_config})" + if step_config is not None + else "" + ), ) use_fused_attention = False fused_attention_backend = None @@ -2468,7 +2477,7 @@ def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_d FP8 attention (O, and dQ/dK/dV under current/mxfp8). It is only consulted when `qkv_dtype` itself is FP8. """ q_format = get_qkv_format(qkv_layout)[1] - eff_qkv_layout = qkv_layout # FP16/BF16 + eff_qkv_layout = qkv_layout # FP16/BF16 if recipe is not None: if not recipe.mxfp8(): # Delayed/current scaling @@ -2492,7 +2501,12 @@ def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_d # FP16/BF16: every tensor is in model precision; scaling_mode is a placeholder if recipe is None: return FusedAttnSpec( - tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, ref, ref, ref, ref, None, + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, + ref, + ref, + ref, + ref, + None, **layout_kwargs, ) @@ -2502,7 +2516,12 @@ def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_d # MXFP8 block scaling: Q/K/V/dO are in MXFP8; O/dQ/dK/dV stay in model precision if recipe.mxfp8(): return FusedAttnSpec( - tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, fprop_fp8, ref, grad_fp8, ref, "bhsd", + tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, + fprop_fp8, + ref, + grad_fp8, + ref, + "bhsd", **layout_kwargs, ) @@ -2511,14 +2530,22 @@ def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_d if recipe.float8_current_scaling(): return FusedAttnSpec( tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, - fprop_fp8, ref if cs_o_in_f16 else fprop_fp8, grad_fp8, ref, None, + fprop_fp8, + ref if cs_o_in_f16 else fprop_fp8, + grad_fp8, + ref, + None, **layout_kwargs, ) # FP8 delayed scaling: Q/K/V/O are in FP8 (e.g. E4M3); dO/dQ/dK/dV in FP8 (e.g. E5M2) return FusedAttnSpec( tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, - fprop_fp8, fprop_fp8, grad_fp8, grad_fp8, None, + fprop_fp8, + fprop_fp8, + grad_fp8, + grad_fp8, + None, **layout_kwargs, ) From cf3265ee7b09a542748a29afb41ed7f13a18a7c6 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 20 Jul 2026 08:08:37 -0700 Subject: [PATCH 35/63] fused-attn: share graph cache across threads via mutex, require cuDNN 9.11, guard zero-size batch/token quantization, add graph-cache debug tracing Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- README.rst | 2 +- docs/installation.rst | 2 +- transformer_engine/common/CMakeLists.txt | 13 +- .../fused_attn_f16_arbitrary_seqlen.cu | 132 +++++++--- .../common/fused_attn/fused_attn_fp8.cu | 112 ++++++--- .../common/fused_attn/graph_debug.h | 233 +++++++++++++++++- transformer_engine/common/fused_attn/utils.cu | 2 + .../dot_product_attention/graph_debug.py | 70 ++++++ .../attention/dot_product_attention/utils.py | 7 +- .../pytorch/cpp_extensions/fused_attn.py | 12 + 10 files changed, 508 insertions(+), 77 deletions(-) create mode 100644 transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py diff --git a/README.rst b/README.rst index 859879fbc2..31bc98e474 100644 --- a/README.rst +++ b/README.rst @@ -160,7 +160,7 @@ System Requirements * **Software:** * CUDA: 12.1+ (Hopper/Ada/Ampere), 12.8+ (Blackwell) with compatible NVIDIA drivers - * cuDNN: 9.3+ + * cuDNN: 9.11+ * Compiler: GCC 9+ or Clang 10+ with C++17 support * Python: 3.12 recommended diff --git a/docs/installation.rst b/docs/installation.rst index cc48a0adac..0271af7fcc 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -14,7 +14,7 @@ Prerequisites 1. Linux x86_64 2. `CUDA 12.1+ (12.8+ for Blackwell support) `__ 3. |driver link|_ supporting CUDA 12.1 or later. -4. `cuDNN 9.3 `__ or later. +4. `cuDNN 9.11 `__ or later. If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within `CUDA_HOME`, set diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 57cdd385b2..ba86420ef5 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -13,8 +13,17 @@ if (CMAKE_BUILD_TYPE STREQUAL "Debug") endif() # Hide non-necessary symbols in shared object. -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") +# [GRAPH-DEBUG] -DNVTE_GRAPH_DEBUG_SYMBOLS=ON keeps + exports internal symbols so +# backtrace_symbols() in fused_attn/graph_debug.h can name fused-attn frames. Remove after +# verification (revert to the two unconditional --version-script lines). +option(NVTE_GRAPH_DEBUG_SYMBOLS "Export all symbols for readable backtraces" OFF) +if (NOT NVTE_GRAPH_DEBUG_SYMBOLS) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") +else() + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -rdynamic -Wl,--export-dynamic") + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -rdynamic -Wl,--export-dynamic") +endif() # Transformer Engine library project(transformer_engine LANGUAGES CUDA CXX) 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 bb593512f5..460cd234b2 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 @@ -10,6 +10,7 @@ #include #include +#include // [SHARED-CACHE] #include #include "../common.h" @@ -182,23 +183,40 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr>; // dropout_offset using CacheType = std::map; - static thread_local CacheType sdpa_f16_fprop_cache; + // [SHARED-CACHE] Process-wide graph cache (was `static thread_local`) so a compiled graph + // is reused across threads instead of rebuilt per thread. Safe because cuDNN >= 9.0 allows + // concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe + // execute(); the static_asserts below fail the build loudly on an older toolkit. + static_assert(CUDNN_VERSION >= 91100, + "[SHARED-CACHE] shared fused-attn graph cache requires cuDNN >= 9.11 " + "(TE minimum supported cuDNN version)"); + static_assert(CUDNN_FRONTEND_VERSION >= 12500, + "[SHARED-CACHE] shared fused-attn graph cache requires cudnn-frontend >= 1.25.0"); + static CacheType sdpa_f16_fprop_cache; + static std::mutex sdpa_f16_fprop_cache_mutex; // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] - fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] - sm_arch_ != 120) { // [GRAPH-DEBUG] - fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building + // so concurrent first-misses on different keys build in parallel. graph->execute() runs + // unlocked after get_graph() returns; built graphs are shared across threads. + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } + fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] + sm_arch_ != 120) { // [GRAPH-DEBUG] + fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] "fwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), - /*legacy=*/!use_cu_seqlens_directly); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] - auto graph = it->second; - return graph; + /*legacy=*/!use_cu_seqlens_directly); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -464,20 +482,32 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); + GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] + GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); + GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); + GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); fused_attn_graph_debug::note_fwd_build(); // [GRAPH-DEBUG] - - return return_tuple; + if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] + std::vector serialized_graph; // [GRAPH-DEBUG] + if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] + fused_attn_graph_debug::note_graph_size("fwd", serialized_graph.size()); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, + // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + { + std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_cache_size("fwd", cache.size()); // [GRAPH-DEBUG] + return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + } }; auto [mha_graph, Q, K, V, attn_scale, O, S1, S2, bias, softmax_offset, seq_q, seq_kv, @@ -731,24 +761,32 @@ void fused_attn_arbitrary_seqlen_bwd_impl( std::shared_ptr>; // dropout_offset using CacheType = std::map; - static thread_local CacheType sdpa_f16_bprop_cache; + static CacheType sdpa_f16_bprop_cache; // [SHARED-CACHE] process-wide (was thread_local) + static std::mutex sdpa_f16_bprop_cache_mutex; // [SHARED-CACHE] // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] - fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] - sm_arch_ != 120) { // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building + // so concurrent first-misses on different keys build in parallel. graph->execute() runs + // unlocked after get_graph() returns; built graphs are shared across threads. + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } + fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] + sm_arch_ != 120) { // [GRAPH-DEBUG] // The backward impl has no cu_seqlens-directly path; it always buckets the batch. fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] "bwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), - /*legacy=*/true); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] - auto graph = it->second; - return graph; + /*legacy=*/true); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -986,19 +1024,31 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); + GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] + GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); + GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); + GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); fused_attn_graph_debug::note_bwd_build(); // [GRAPH-DEBUG] - - return return_tuple; + if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] + std::vector serialized_graph; // [GRAPH-DEBUG] + if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] + fused_attn_graph_debug::note_graph_size("bwd", serialized_graph.size()); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, + // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + { + std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_cache_size("bwd", cache.size()); // [GRAPH-DEBUG] + return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + } }; auto [mha_graph, q, k, v, o, dO, stats, attn_scale, dQ, dK, dV, bias, dBias, softmax_offset, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 5163a3601a..5eddd8fa75 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -4,6 +4,9 @@ * See LICENSE for license information. ************************************************************************/ +#include // [SHARED-CACHE] +#include // [GRAPH-DEBUG] serialized-size probe + #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" @@ -128,17 +131,34 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de std::shared_ptr>; // dropout_offset using CacheType = std::map; - static thread_local CacheType sdpa_fp8_fprop_cache; + // [SHARED-CACHE] Process-wide graph cache (was `static thread_local`) so a compiled graph + // is reused across threads instead of rebuilt per thread. Safe because cuDNN >= 9.0 allows + // concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe + // execute(); the static_asserts below fail the build loudly on an older toolkit. + static_assert(CUDNN_VERSION >= 91100, + "[SHARED-CACHE] shared fused-attn graph cache requires cuDNN >= 9.11 " + "(TE minimum supported cuDNN version)"); + static_assert(CUDNN_FRONTEND_VERSION >= 12500, + "[SHARED-CACHE] shared fused-attn graph cache requires cudnn-frontend >= 1.25.0"); + static CacheType sdpa_fp8_fprop_cache; + static std::mutex sdpa_fp8_fprop_cache_mutex; // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building + // so concurrent first-misses on different keys build in parallel. graph->execute() runs + // unlocked after get_graph() returns; built graphs are shared across threads. + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] - auto graph = it->second; - return graph; + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -388,18 +408,30 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); + GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] + GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); + GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); + GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); fused_attn_graph_debug::note_fwd_build(); // [GRAPH-DEBUG] - - return return_tuple; + if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] + std::vector serialized_graph; // [GRAPH-DEBUG] + if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] + fused_attn_graph_debug::note_graph_size("fwd", serialized_graph.size()); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, + // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + { + std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_cache_size("fwd", cache.size()); // [GRAPH-DEBUG] + return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + } }; auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, @@ -611,17 +643,25 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr>; // dropout_offset using CacheType = std::map; - static thread_local CacheType sdpa_fp8_bprop_cache; + static CacheType sdpa_fp8_bprop_cache; // [SHARED-CACHE] process-wide (was thread_local) + static std::mutex sdpa_fp8_bprop_cache_mutex; // [SHARED-CACHE] // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - bool cache_hit = (it != cache.end()); // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building + // so concurrent first-misses on different keys build in parallel. graph->execute() runs + // unlocked after get_graph() returns; built graphs are shared across threads. + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] - auto graph = it->second; - return graph; + if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -1000,19 +1040,31 @@ void fused_attn_fp8_bwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); + GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] + GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); + GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); + GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); fused_attn_graph_debug::note_bwd_build(); // [GRAPH-DEBUG] - - return return_tuple; + if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] + std::vector serialized_graph; // [GRAPH-DEBUG] + if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] + fused_attn_graph_debug::note_graph_size("bwd", serialized_graph.size()); // [GRAPH-DEBUG] + } // [GRAPH-DEBUG] + // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, + // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + { + std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + fused_attn_graph_debug::note_cache_size("bwd", cache.size()); // [GRAPH-DEBUG] + return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + } }; auto [mha_graph, Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, diff --git a/transformer_engine/common/fused_attn/graph_debug.h b/transformer_engine/common/fused_attn/graph_debug.h index 8a69b85e78..258a4556c4 100644 --- a/transformer_engine/common/fused_attn/graph_debug.h +++ b/transformer_engine/common/fused_attn/graph_debug.h @@ -19,7 +19,9 @@ // which impl path (bucketed batch vs. real batch) the graph was built for. // - A "SUMMARY" line with final build/exec totals is printed at process exit, followed by a // "THD-PATH" line with per-path lookup/build totals (low builds/lookups on the legacy path -// means batch bucketing is collapsing distinct batch sizes onto shared graphs). +// means batch bucketing is collapsing distinct batch sizes onto shared graphs), and one +// "STAGE " line per FE build stage (validate ... build_plans) with total CPU/wall +// time and call count -- on `main` these were static boolean checks (~0 cost). // // Separately, force every lookup to miss (never reuse a cached graph) with: // export NVTE_FUSED_ATTN_DISABLE_CACHE=1 @@ -36,11 +38,28 @@ #ifndef TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ #define TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ +#include #include +#include #include #include #include +#include +#include +#include #include +#include + +// [GRAPH-DEBUG] Backtrace printing needs glibc's and libstdc++'s +// (demangling). Gate on availability so non-glibc toolchains still build; dump_backtrace() +// becomes a no-op there. +#if defined(__has_include) +#if __has_include() && __has_include() +#define NVTE_FUSED_ATTN_GRAPH_DEBUG_HAVE_BACKTRACE 1 +#include +#include +#endif +#endif #include "config_and_params.h" // [GRAPH-DEBUG] for FusedAttnConfig field dump @@ -126,12 +145,100 @@ inline void dump_thd_summary() { std::fflush(stderr); } +// [GRAPH-DEBUG] Host-memory footprint of cached graphs, split fwd/bwd (index 0/1). +// serialized bytes: size of fe::graph::Graph::serialize() output -- a proxy for the host +// memory one built graph holds (its plan / engine config / tensor metadata). Summed over +// builds; avg = sum / count gives the per-graph host cost. +// cache entries: high-water number of live graphs in the shared std::map (one graph per key). +// Device memory (workspace) is separate and sized per execute(), not held by the cached graph. +inline int pass_index(const char *pass) { return (pass[0] == 'b') ? 1 : 0; } // "bwd" -> 1 + +inline std::atomic &serial_bytes(int i) { + static std::array, 2> v{}; + return v[i]; +} +inline std::atomic &serial_count(int i) { + static std::array, 2> v{}; + return v[i]; +} +inline std::atomic &cache_entries(int i) { + static std::array, 2> v{}; + return v[i]; +} + +inline void dump_memory_summary() { + for (int i = 0; i < 2; ++i) { + const char *pass = (i == 0) ? "fwd" : "bwd"; + uint64_t cnt = serial_count(i).load(); + uint64_t bytes = serial_bytes(i).load(); + uint64_t entries = cache_entries(i).load(); + double total_kb = static_cast(bytes) / 1024.0; + double avg_kb = cnt ? total_kb / static_cast(cnt) : 0.0; + std::fprintf( + stderr, + "[GRAPH-DEBUG] MEMORY %-3s | cache entries=%llu | serialized graphs=%llu total=%.1f KB (avg %.1f KB)\n", + pass, static_cast(entries), static_cast(cnt), + total_kb, avg_kb); + } + std::fflush(stderr); +} + +// [GRAPH-DEBUG] Per-stage CPU/wall time for the FE build pipeline (validate ... build_plans). +// On `main` these were static boolean checks (~0 cost); this quantifies the added cost. +enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; + +inline const char *stage_name(BuildStage s) { + switch (s) { + case BuildStage::Validate: + return "validate"; + case BuildStage::BuildOpGraph: + return "build_operation_graph"; + case BuildStage::CreatePlans: + return "create_execution_plans"; + case BuildStage::CheckSupport: + return "check_support"; + case BuildStage::BuildPlans: + return "build_plans"; + default: + return "?"; + } +} + +inline std::atomic &stage_calls(BuildStage s) { + static std::array, static_cast(BuildStage::kCount)> v{}; + return v[static_cast(s)]; +} +inline std::atomic &stage_cpu_ns(BuildStage s) { + static std::array, static_cast(BuildStage::kCount)> v{}; + return v[static_cast(s)]; +} +inline std::atomic &stage_wall_ns(BuildStage s) { + static std::array, static_cast(BuildStage::kCount)> v{}; + return v[static_cast(s)]; +} + +inline void dump_stage_summary() { + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + BuildStage s = static_cast(i); + uint64_t n = stage_calls(s).load(); + if (n == 0) continue; + double cpu_ms = static_cast(stage_cpu_ns(s).load()) / 1e6; + double wall_ms = static_cast(stage_wall_ns(s).load()) / 1e6; + std::fprintf(stderr, + "[GRAPH-DEBUG] STAGE %-22s | calls=%llu | cpu=%.1f ms (avg %.3f ms) | wall=%.1f ms\n", + stage_name(s), static_cast(n), cpu_ms, cpu_ms / n, wall_ms); + } + std::fflush(stderr); +} + inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (enabled()) { dump("SUMMARY"); dump_thd_summary(); + dump_stage_summary(); + dump_memory_summary(); } }); return true; @@ -162,6 +269,26 @@ inline void note_bwd_exec() { bwd_exec().fetch_add(1); } +// [GRAPH-DEBUG] Record the serialized size (host-memory proxy) of one freshly built graph. +// Call only after a successful serialize() so the average reflects real graphs. +inline void note_graph_size(const char *pass, size_t serialized_bytes) { + if (!enabled()) return; + register_summary_once(); + int i = pass_index(pass); + serial_bytes(i).fetch_add(serialized_bytes); + serial_count(i).fetch_add(1); +} + +// [GRAPH-DEBUG] Record the current shared-cache entry count (kept as a high-water mark). +inline void note_cache_size(const char *pass, size_t entries) { + if (!enabled()) return; + register_summary_once(); + int i = pass_index(pass); + uint64_t prev = cache_entries(i).load(); + while (entries > prev && !cache_entries(i).compare_exchange_weak(prev, entries)) { + } +} + // Returns true when the graph cache should be bypassed (every lookup treated as a // miss so a fresh graph is built each call). Gated by NVTE_FUSED_ATTN_DISABLE_CACHE. inline bool cache_disabled() { @@ -172,6 +299,67 @@ inline bool cache_disabled() { return off; } +// [GRAPH-DEBUG] Opt-in C++ backtrace printing next to each cache lookup. Kept separate from the +// main NVTE_FUSED_ATTN_GRAPH_DEBUG switch because a full stack per lookup is very verbose; enable +// with NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE=1 (the main switch must also be on). +inline bool backtrace_enabled() { + static const bool on = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return on; +} + +// [GRAPH-DEBUG] Frames to print per lookup (override with NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE_DEPTH). +inline int backtrace_depth() { + static const int depth = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE_DEPTH"); + int d = (e != nullptr && e[0] != '\0') ? std::atoi(e) : 24; + if (d < 1) d = 1; + if (d > 128) d = 128; + return d; + }(); + return depth; +} + +// [GRAPH-DEBUG] Print a symbolized (and, where possible, demangled) C++ backtrace, one frame per +// line, each tagged so the frames group visually under the HIT/MISS line they belong to. `skip` +// drops the top frames that are just this instrumentation (dump_backtrace + its caller). For +// readable function names the library must be built/linked with -rdynamic (or -g); otherwise +// non-exported frames show as "(+0x)". +inline void dump_backtrace(const char *tag, int skip = 2) { + if (!backtrace_enabled()) return; +#if defined(NVTE_FUSED_ATTN_GRAPH_DEBUG_HAVE_BACKTRACE) + const int max_frames = backtrace_depth() + skip; + std::vector frames(static_cast(max_frames)); + int n = ::backtrace(frames.data(), max_frames); + if (n <= skip) return; + char **symbols = ::backtrace_symbols(frames.data(), n); + if (symbols == nullptr) return; + for (int i = skip; i < n; ++i) { + // glibc format: "(+0x) [0x]"; demangle the "" span. + std::string line = symbols[i]; + char *open = std::strchr(symbols[i], '('); + char *plus = open ? std::strchr(open, '+') : nullptr; + if (open != nullptr && plus != nullptr && plus > open + 1) { + std::string mangled(open + 1, plus); + int status = 0; + char *demangled = abi::__cxa_demangle(mangled.c_str(), nullptr, nullptr, &status); + if (status == 0 && demangled != nullptr) { + line = std::string(symbols[i], open + 1) + demangled + plus; + std::free(demangled); + } + } + std::fprintf(stderr, "[GRAPH-DEBUG] bt[%-4s] #%02d %s\n", tag, i - skip, line.c_str()); + } + std::fflush(stderr); + std::free(symbols); +#else + (void)tag; + (void)skip; +#endif +} + // Logs one graph-cache lookup with its outcome (HIT/MISS) and the *real* (pre- // normalization) config fields. A std::map HIT means the two configs compare equal // under operator<, so the field that actually distinguishes a wrongly-reused graph @@ -216,6 +404,7 @@ inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), static_cast(c.bias_seqlen_kv)); std::fflush(stderr); + dump_backtrace(hit ? "HIT" : "MISS"); // [GRAPH-DEBUG] frames for this fwd/bwd lookup } // Records, for one THD (ragged) cache lookup, which impl path the graph was built for -- @@ -237,9 +426,51 @@ inline void note_thd_lookup(const char *pass, bool hit, bool built, bool legacy) hit ? "HIT" : "MISS", thread_seq_id(), legacy ? "legacy" : "direct", (hit && built) ? " [cache-disabled->rebuild]" : ""); std::fflush(stderr); + dump_backtrace(hit ? "HIT" : "MISS"); // [GRAPH-DEBUG] frames for this THD (ragged) lookup } +// [GRAPH-DEBUG] Thread-CPU clock (excludes time blocked on locks / GPU sync), in nanoseconds. +inline uint64_t cpu_now_ns() { + timespec ts; + clock_gettime(CLOCK_THREAD_CPUTIME_ID, &ts); + return static_cast(ts.tv_sec) * 1000000000ull + static_cast(ts.tv_nsec); +} + +// [GRAPH-DEBUG] RAII timer: records wall + thread-CPU time for one FE build stage. Zero cost +// (only an enabled() bool read) when NVTE_FUSED_ATTN_GRAPH_DEBUG is unset. The destructor records +// even on early return / thrown NVTE_CHECK, so it is safe to wrap the checked FE calls. +struct ScopedStageTimer { + BuildStage stage; + bool on; + std::chrono::steady_clock::time_point w0; + uint64_t c0{0}; + explicit ScopedStageTimer(BuildStage s) : stage(s), on(enabled()) { + if (!on) return; + register_summary_once(); + c0 = cpu_now_ns(); + w0 = std::chrono::steady_clock::now(); + } + ~ScopedStageTimer() { + if (!on) return; + uint64_t cpu = cpu_now_ns() - c0; + uint64_t wall = static_cast( + std::chrono::duration_cast(std::chrono::steady_clock::now() - w0) + .count()); + stage_cpu_ns(stage).fetch_add(cpu); + stage_wall_ns(stage).fetch_add(wall); + stage_calls(stage).fetch_add(1); + } +}; + } // namespace fused_attn_graph_debug } // namespace transformer_engine +// [GRAPH-DEBUG] Wrap a single (possibly NVTE_CHECK_*-guarded) FE call to time it under `stage`. +#define GRAPH_DEBUG_TIME_STAGE(stage, expr) \ + do { \ + ::transformer_engine::fused_attn_graph_debug::ScopedStageTimer _gd_stage_timer( \ + ::transformer_engine::fused_attn_graph_debug::BuildStage::stage); \ + expr; \ + } while (0) + #endif // TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 77122c6424..875ccdbe72 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -510,6 +510,7 @@ DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_at // quantize batch size size_t get_max_batch_size(size_t batch_size) { + if (batch_size == 0) return 0; // guard: log2(0) = -inf, casting to size_t is UB size_t max_b = batch_size; size_t log2_b = ceil(log2(batch_size)); // batch size is expected to be 10s-100s @@ -528,6 +529,7 @@ size_t get_max_batch_size(size_t batch_size) { // quantize token count size_t get_max_tokens(size_t num_tokens) { + if (num_tokens == 0) return 0; // guard: log2(0) = -inf, casting to size_t is UB // token count is expected to be 1k's-100k's // t = 0, ..., 1024 -> max_t = 1024 // t = 1025, ..., 32k -> max_t = next power of 2 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py b/transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py new file mode 100644 index 0000000000..b69839f74c --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py @@ -0,0 +1,70 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# ============================================================================ +# [GRAPH-DEBUG] TEMPORARY DEBUG INSTRUMENTATION -- REMOVE AFTER VERIFICATION. +# +# Python-side companion to the C++ instrumentation in +# common/fused_attn/graph_debug.h. Prints the Python call stack that leads into +# each fused-attention backend query / forward / backward call, so the Python +# frames interleave (on stderr) just above the C++ "[GRAPH-DEBUG] fwd/bwd HIT|MISS" +# lines they trigger. This makes it possible to attribute each cuDNN graph-cache +# lookup to the exact Python caller (availability probe vs. module backend +# re-selection vs. actual fwd/bwd execution). +# +# Enable with the SAME switch as the C++ side: +# export NVTE_FUSED_ATTN_GRAPH_DEBUG=1 +# Optionally cap the number of printed frames (default 12): +# export NVTE_FUSED_ATTN_GRAPH_DEBUG_PY_DEPTH= +# +# To remove all of this instrumentation later: +# 1. Delete this file (graph_debug.py). +# 2. Remove every line tagged with the "[GRAPH-DEBUG]" marker in: +# - attention/dot_product_attention/utils.py +# - cpp_extensions/fused_attn.py +# ============================================================================ + +import os +import sys +import threading +import traceback + +_enabled = None +_depth = None + + +def enabled(): + """True when NVTE_FUSED_ATTN_GRAPH_DEBUG is set (same switch as the C++ side).""" + global _enabled + if _enabled is None: + val = os.getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG", "") + _enabled = val not in ("", "0") + return _enabled + + +def _depth_val(): + global _depth + if _depth is None: + val = os.getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG_PY_DEPTH", "") + try: + _depth = int(val) if val else 12 + except ValueError: + _depth = 12 + _depth = max(1, min(_depth, 128)) + return _depth + + +def pytrace(tag): + """Print a compact Python call stack to stderr, tagged so it groups with the C++ + [GRAPH-DEBUG] frames that follow. No-op unless NVTE_FUSED_ATTN_GRAPH_DEBUG is set.""" + if not enabled(): + return + # Drop this frame (pytrace itself); show the most recent frames, oldest first. + frames = traceback.extract_stack()[:-1][-_depth_val() :] + out = sys.stderr + out.write(f"[GRAPH-DEBUG-PY] {tag} | tid={threading.get_ident()}\n") + for fr in frames: + code = f" -> {fr.line}" if fr.line else "" + out.write(f"[GRAPH-DEBUG-PY] {fr.filename}:{fr.lineno} {fr.name}(){code}\n") + out.flush() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index c3b9c1da1b..76685874aa 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -316,7 +316,7 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 - softmax_scale: float = 0.0 + softmax_scale: float = 1.0 fp8_output: bool = False checkpoint_core_attention: bool = False has_score_mod: bool = False @@ -426,6 +426,11 @@ def get_attention_backend( All available backends that could support the provided input. A list of Booleans in the form of [use_flash_attention, use_fused_attention, use_unfused_attention]. """ + # [GRAPH-DEBUG] Trace the Python caller that triggers a fused-attn backend query (maps to the + # C++ support-check "fwd/bwd HIT|MISS" lines from is_supported_f16_*). Remove after verification. + from transformer_engine.pytorch.attention.dot_product_attention import graph_debug + + graph_debug.pytrace("get_attention_backend") # NOTE: As part of refactoring attention.py, populating the _attention_backends cache in attention # is no longer performed at the end of get_attention_backend(), but the responsibility of doing so # is shifted over to the caller of this function diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9c22c56bd1..3e68eab85b 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -303,6 +303,12 @@ def fused_attn_fwd( else: raise ValueError(f"Unsupported backend {fused_attention_backend}") + # [GRAPH-DEBUG] Trace the Python caller of the actual fwd kernel (maps to the C++ execution + # "fwd HIT|MISS" line + note_fwd_exec). Remove after verification. + from transformer_engine.pytorch.attention.dot_product_attention import graph_debug + + graph_debug.pytrace("fused_attn_fwd (execute)") + # execute kernel output_tensors = tex.fused_attn_fwd( max_seqlen_q, @@ -553,6 +559,12 @@ def fused_attn_bwd( f" for backend={fused_attention_backend}." ) + # [GRAPH-DEBUG] Trace the Python caller of the actual bwd kernel (maps to the C++ execution + # "bwd HIT|MISS" line + note_bwd_exec). Remove after verification. + from transformer_engine.pytorch.attention.dot_product_attention import graph_debug + + graph_debug.pytrace("fused_attn_bwd (execute)") + output_tensors = tex.fused_attn_bwd( max_seqlen_q, max_seqlen_kv, From cda01d7142a3bca27c3f919ae732ccba5c24d39f Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 06:31:37 -0700 Subject: [PATCH 36/63] Cache attention-backend selection keyed on (NVTE_* env, attention_params), drop temporary graph-debug instrumentation, revert min cudnn version update Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- README.rst | 2 +- docs/envvars.rst | 18 +- docs/examples/attention/attention.ipynb | 1232 ++++++++--------- docs/installation.rst | 2 +- tests/pytorch/attention/test_attention.py | 17 - tests/pytorch/utils.py | 14 +- transformer_engine/common/CMakeLists.txt | 13 +- .../common/fused_attn/config_and_params.cpp | 21 +- .../common/fused_attn/config_and_params.h | 2 + .../common/fused_attn/fused_attn.cpp | 9 +- .../fused_attn_f16_arbitrary_seqlen.cu | 127 +- .../fused_attn_f16_arbitrary_seqlen.h | 14 +- .../common/fused_attn/fused_attn_fp8.cu | 77 +- .../common/fused_attn/fused_attn_fp8.h | 13 +- .../common/fused_attn/graph_cache_debug.h | 274 ++++ .../common/fused_attn/graph_debug.h | 476 ------- transformer_engine/common/fused_attn/utils.cu | 88 -- transformer_engine/common/fused_attn/utils.h | 36 +- .../dot_product_attention.py | 145 +- .../dot_product_attention/graph_debug.py | 70 - .../attention/dot_product_attention/utils.py | 5 - .../pytorch/cpp_extensions/fused_attn.py | 12 - 22 files changed, 1132 insertions(+), 1535 deletions(-) create mode 100644 transformer_engine/common/fused_attn/graph_cache_debug.h delete mode 100644 transformer_engine/common/fused_attn/graph_debug.h delete mode 100644 transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py diff --git a/README.rst b/README.rst index 31bc98e474..859879fbc2 100644 --- a/README.rst +++ b/README.rst @@ -160,7 +160,7 @@ System Requirements * **Software:** * CUDA: 12.1+ (Hopper/Ada/Ampere), 12.8+ (Blackwell) with compatible NVIDIA drivers - * cuDNN: 9.11+ + * cuDNN: 9.3+ * Compiler: GCC 9+ or Clang 10+ with C++17 support * Python: 3.12 recommended diff --git a/docs/envvars.rst b/docs/envvars.rst index b3765a06bd..e543975f59 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -171,18 +171,24 @@ backend-selection overview. :Default: ``1`` :Description: Enable or disable UnfusedDotProductAttention backend (native PyTorch). When set to ``0``, UnfusedDotProductAttention will not be used. -.. envvar:: NVTE_FUSED_ATTN_BACKEND - - :Type: ``int`` (1 or 2) - :Default: Auto-selected - :Description: Request a cuDNN FusedAttention backend when that request is supported by the active fused-attention path. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. BF16/FP16 attention uses sub-backend ``1`` when eligible. FP8 attention uses sub-backend ``2`` when FP8 DPA is enabled and supported by the architecture, cuDNN version, and input configuration. - .. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD :Type: ``int`` (0 or 1) :Default: ``0`` :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. +.. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable diagnostic logging for the cuDNN FusedAttention graph cache. When set to ``1``, prints to stderr (prefixed ``[FUSED-ATTN-CACHE]``) a per-lookup HIT/MISS line with the full graph-cache key, a BUILD line whenever a new graph is constructed, and a SUMMARY of graph builds vs. executions at process exit. Useful for diagnosing redundant graph rebuilds or stale-cache reuse. Has negligible overhead when unset. + +.. envvar:: NVTE_FUSED_ATTN_DISABLE_CACHE + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Bypass the cuDNN FusedAttention graph cache, rebuilding a fresh graph on every forward/backward call. Intended for debugging stale-cache reuse only: it forces expensive graph recompilation on every call and must not be used in production. If a run that fails with the cache enabled passes with it disabled, the bug is stale-cache reuse (an incomplete cache key). Pairs with :envvar:`NVTE_FUSED_ATTN_CACHE_DEBUG` for inspecting each rebuild. + .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO :Type: ``int`` (0 or 1) diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 989661b543..6c868518ec 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -1,628 +1,618 @@ { - "cells": [ - { - "cell_type": "markdown", - "id": "040f466a", - "metadata": {}, - "source": [ - "# Attention Is All You Need!\n", - "\n", - "The core idea behind Transformer models is the attention mechanism [[1]](https://arxiv.org/abs/1706.03762). It identifies the correlation between words, selects the most important parts of the sentence to focus on, and captures meaningful patterns and dependencies in the data. Figure 1 shows a typical attention mechanism, where pre-softmax operations can be a combination of scaling, bias and masking while the post-softmax operation is often just dropout.\n", - "\n", - "
\n", - "\n", - "
Figure 1: Dot product attention.
\n", - "
\n", - "\n", - "[Transformer Engine](https://github.com/NVIDIA/TransformerEngine.git) supports the calculation of dot product attention in two frameworks, [PyTorch](https://github.com/pytorch/pytorch) and [JAX](https://github.com/google/jax). The API for each framework is\n", - "\n", - "- [transformer_engine.pytorch.DotProductAttention](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention)\n", - "- [transformer_engine.jax.flax.DotProductAttention](../../api/jax.rst#transformer_engine.jax.flax.DotProductAttention)" - ] - }, - { - "cell_type": "markdown", - "id": "89a7d849", - "metadata": {}, - "source": [ - "## 1. Attention Backends\n", - "\n", - "Transformer Engine provides multiple attention backends for each supported framework. The framework-native backends provide a robust baseline, while the fused, GPU-optimized implementations offer more performance. For example, the flash-attention and cuDNN attention backends in PyTorch. The framework-native backends are often named with \"unfused\", while the more optimized backends are \"fused\" or \"flash\".\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
FrameworkBackend (Module Name)Module Location
PyTorchcuDNN attention (`FusedAttention`) [transformer_engine.pytorch.attention](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py)
flash-attention (`FlashAttention`)
\n", - " PyTorch-native attention (`UnfusedDotProductAttention`)\n", - "
JAXcuDNN attention (`_FusedDotProductAttention`)[transformer_engine.jax.flax.transformer](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/flax/transformer.py)
JAX-native attention (`_UnfusedDotProductAttention`)
" - ] - }, - { - "cell_type": "markdown", - "id": "c90a2573", - "metadata": {}, - "source": [ - "### 1.1 Flash vs. Non-Flash\n", - "\n", - "The attention calculation has quadratic computational and memory complexities to the sequence length. Its runtime and memory requirements quadruple, when the sequence length doubles. This presents a significant challenge to scale Transformer models up for longer contexts, in order to achieve higher model quality.\n", - "\n", - "Compared to the standard, non-flash algorithm, the flash algorithm [[2]](https://arxiv.org/abs/2205.14135) was proposed to reduce the memory scaling to linear and improve the computational efficiency through optimized memory accesses. It employs the following two distinctive techniques.\n", - "\n", - "- **Tiling:** The non-flash algorithm tries to process the query, key, value tensors in one single step, requiring large amounts of global memory and incurring high volumes of reads/writes between global memory and shared memory. The flash algorithm decomposes the input into several tiles, based on the available shared memory and register size, and it computes the softmax one tile at a time.\n", - "\n", - "- **Recomputation:** The non-flash algorithm stores the softmax matrix (quadratic to sequence length) to global memory for the backward pass, while the flash algorithm only saves the softmax normalization factors (linear to sequence length). This reduces the amount of memory required as well as the bandwidth utilization between global memory and shared memory. Even though there is extra computation incurred in order to recalculate the attention in the backward pass, the bandwidth savings still provide significant improvement in efficiency.\n", - "\n", - "
\n", - "Note: \n", - " \n", - "Transformer Engine's flash-attention backend, available in PyTorch, and cuDNN attention backend (sub-backends 1 and 2), available in PyTorch and JAX, are both based on the flash algorithm.\n", - "
\n" - ] - }, - { - "cell_type": "markdown", - "id": "b5ce567d", - "metadata": {}, - "source": [ - "### 1.2 flash-attention\n", - "\n", - "The flash-attention backend, available only in PyTorch, is a module wrapped around the public `flash-attn` package [[3]](https://github.com/Dao-AILab/flash-attention). \n", - "\n", - "The flash-attention backend supports `flash-attn`'s features as well as a few extra functionalities to facilitate the use of `flash-attn`, such as converting the `attention_mask` to cumulative sequence lengths `cu_seqlens` for `padding` mask use cases. Please see `transformer_engine.pytorch.attention.FlashAttention` for details.\n", - "\n", - "The `flash-attn` dependency is regularly updated in Transformer Engine. As of v2.0, Transformer Engine supports `flash-attn` 2.0.6+ (see [setup.py](https://github.com/NVIDIA/TransformerEngine/blob/main/setup.py)).\n", - "\n", - "To understand `flash-attn`'s performance, please refer to their benchmarks [here](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#performance).\n", - "\n", - "### 1.3 cuDNN Attention\n", - "\n", - "The cuDNN attention backend, available in PyTorch and JAX, offers another high-performance solution to the attention calculation. It requires [cuDNN](https://developer.nvidia.com/cudnn) to run, and has several sub-backends to support the different precisions and sequence lengths.\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
Sub-BackendAlgorithmPrecisionSequence LengthArchitectureAdditional info
1FlashBF16/FP16 Any sm80+ [cuDNN](https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#fused-flash-attention-fprop),\n", - " [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention)\n", - "
2FlashFP8 cuDNN pre-9.0: ≤512 cuDNN pre-9.0: sm90
cuDNN 9.0+: Any cuDNN 9.0+: sm90+ cuDNN 9.0+: [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention-fp8)\n", - "
\n", - "\n", - "The cuDNN attention backend and flash-attention backend have several notable differences. As of Transformer Engine 2.0, cuDNN 9.3 and `flash-attn` 2.4.2,\n", - "\n", - "- flash-attention only supports the PyTorch framework while cuDNN attention supports PyTorch and JAX.\n", - "- flash-attention supports BF16, FP16 precisions while cuDNN attention also supports FP8 (through its sub-backend 2).\n", - "- flash-attention supports `bshd`, `thd` input formats, without any transposes, and `sbhd` format, with transposes, while cuDNN attention supports all three formats without transposes (see Section 3.1 for more details).\n", - "- flash-attention does not support `post_scale_bias`, and cuDNN attention does.\n", - "- flash-attention supports KV-caching and paged attention, and cuDNN attention does not.\n", - "- flash-attention uses bottom right diagonal for `causal` mask in cross attention (see [change log](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#21-change-behavior-of-causal-flag)), and cuDNN attention supports both top left and bottom right.\n", - "- **Sliding window attention (SWA):** flash-attention has SWA(left, right) support for all mask types except top-left causal masks, with or without dropout, and without bias. cuDNN attention supports SWA(left, 0) starting from 9.2 and SWA(left, right) starting from 9.6, without dropout, and with `bias_type=\"no_bias\"`.\n", - "- flash-attention outperforms cuDNN attention on Ampere architectures, and cuDNN attention has 20-50% advantages on Hopper architectures, based on our benchmarks for a number of commonly-used model configurations.\n", - "\n", - "To compare cuDNN attention and flash-attention, users can modify the `model_configs` dictionary in [benchmarks/attention/benchmark_attention.py](https://github.com/NVIDIA/TransformerEngine/blob/main/benchmarks/attention/benchmark_attention.py) to collect performance numbers. The script runs each entry in `model_configs` for `num_iters` times, each time with one forward pass and one backward pass. Both backends are tried, and if one backend does not have support for the specific user input, the runtimes and speedups in the final table would be 0." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "c5b8e3d7", - "metadata": {}, - "outputs": [], - "source": [ - "model_configs = {\n", - " # test: b, h, hg, d, sq, skv, p, mask, bias\n", - " \"test_0\": ModelConfig(2, 16, 16, 64, 512, 512, 0.0, \"no_mask\", \"no_bias\"), # short seq\n", - " \"test_1\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"no_bias\"), # longer seq, mask\n", - " \"test_2\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"post_scale_bias\"), # bias\n", - " \"test_3\": ModelConfig(2, 32, 4, 128, 8192, 8192, 0.0, \"causal\", \"no_bias\"), # GQA\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "50852cb5", - "metadata": {}, - "outputs": [ + "cells": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Device 0: NVIDIA H100 80GB HBM3 GPU, sm90 compute capability, 79.1GB memory\n", - "Running test_0 with cuDNN attention and flash-attention...\n", - "Running test_1 with cuDNN attention and flash-attention...\n", - "Running test_2 with cuDNN attention...\n", - "Running test_3 with cuDNN attention and flash-attention...\n", - "\n", - " cuDNN fwd+bwd (ms) flash-attn fwd+bwd (ms) cuDNN vs flash speedup\n", - "test_0 0.0340 0.0468 1.3786\n", - "test_1 0.3664 0.5850 1.5968\n", - "test_2 0.9332 0.0000 0.0000\n", - "test_3 7.4875 11.8879 1.5877\n" - ] - } - ], - "source": [ - "!cd ../../../benchmarks/attention/ && python benchmark_attention.py" - ] - }, - { - "cell_type": "markdown", - "id": "9a615119", - "metadata": {}, - "source": [ - "## 2. Backend Selection\n", - "\n", - "Given the various attention backends, Transformer Engine first determines which backends are eligible for the provided inputs and runtime environment, then applies a preference order among the eligible backends. Eligibility is affected by user environment variables, GPU architecture, installed `flash-attn` and cuDNN versions, data type and FP8 recipe, QKV layout, training or inference mode, dropout, and other attention features.\n", - "\n", - "In PyTorch, the candidates are FlashAttention (`flash-attn` v2, v3, or v4), FusedAttention (cuDNN sub-backends), and UnfusedDotProductAttention. Users can disable whole backend families with `NVTE_FLASH_ATTN`, `NVTE_FUSED_ATTN`, or `NVTE_UNFUSED_ATTN`. In JAX, Transformer Engine checks whether a cuDNN fused-attention kernel is available when `NVTE_FUSED_ATTN=1`; otherwise it falls back to the JAX-native implementation.\n", - "\n", - "At a high level, the architecture-specific PyTorch selection order is:\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
FrameworkSelection Order
PyTorchsm8x (Ampere/Ada): flash-attention > cuDNN attention > PyTorch-native attention
sm90 (Hopper): cuDNN attention > flash-attention > PyTorch-native attention
sm100/sm120 (Blackwell): cuDNN attention > flash-attention > PyTorch-native attention
cuDNN attention: BF16/FP16 uses sub-backend 1 when eligible; FP8 uses sub-backend 2 when enabled and eligible
JAXcuDNN attention > JAX-native attention
\n", - "\n", - "Within FlashAttention, TE uses the installed implementation that is supported for the architecture and input. FlashAttention 3 is Hopper-only (`sm90`). FlashAttention 4 supports `sm80`, `sm90`, `sm100`, and `sm120`; on Hopper, TE prefers FlashAttention 3 over FlashAttention 4 when both are installed and eligible. On Blackwell, FlashAttention 4 is the Blackwell-specific flash-attention path when installed and eligible, while FlashAttention 2 can still be eligible depending on the installed version and input configuration.\n", - "\n", - "Within cuDNN FusedAttention, TE asks the fused-attention helper which sub-backend is eligible. Sub-backend 1 is the BF16/FP16 flash-based path when available; sub-backend 2 is the FP8 path when FP8 DPA is enabled and the architecture, cuDNN version, and input configuration support it. Hopper supports eligible FP8 DPA through cuDNN sub-backend 2. In the current PyTorch selector, eligible FP8 DPA on Blackwell is an `sm100` path and is disabled on `sm120`.\n", - "\n", - "When all optimized backends are disabled or ineligible, TE falls back to UnfusedDotProductAttention if it is enabled. If no backend is eligible, backend selection returns no backend and the caller raises an error. As we monitor the performance of different backends, the selection logic may change." - ] - }, - { - "cell_type": "markdown", - "id": "e6c0f3f0", - "metadata": {}, - "source": [ - "### 2.1 Debug Information\n", - "\n", - "To find out which backend is being used during runtime, we have the following two debugging flags. Logging is done by using the `logging` package.\n", - "```\n", - "NVTE_DEBUG = 0/1 # disables/enables debugging\n", - "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n", - "```\n", - "
\n", - "Note:\n", - " \n", - "These flags are supported in PyTorch only as of Transformer Engine 2.0. JAX support is expected to be added in the future.\n", - "
" - ] - }, - { - "cell_type": "markdown", - "id": "16660323", - "metadata": {}, - "source": [ - "The example script [example_attention.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/example_attention.py) runs a very basic model with two attention backends, cuDNN attention and flash-attention. Here `NVTE_DEBUG_LEVEL=1` allows us to find out which backend/sub-backend is used in runtime." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "906b8cf1", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Attention Is All You Need!\n", + "\n", + "The core idea behind Transformer models is the attention mechanism [[1]](https://arxiv.org/abs/1706.03762). It identifies the correlation between words, selects the most important parts of the sentence to focus on, and captures meaningful patterns and dependencies in the data. Figure 1 shows a typical attention mechanism, where pre-softmax operations can be a combination of scaling, bias and masking while the post-softmax operation is often just dropout.\n", + "\n", + "
\n", + "\n", + "
Figure 1: Dot product attention.
\n", + "
\n", + "\n", + "[Transformer Engine](https://github.com/NVIDIA/TransformerEngine.git) supports the calculation of dot product attention in two frameworks, [PyTorch](https://github.com/pytorch/pytorch) and [JAX](https://github.com/google/jax). The API for each framework is\n", + "\n", + "- [transformer_engine.pytorch.DotProductAttention](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention)\n", + "- [transformer_engine.jax.flax.DotProductAttention](../../api/jax.rst#transformer_engine.jax.flax.DotProductAttention)" + ], + "id": "040f466a" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Run cuDNN attention...\n", - "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", - "\n", - "Run flash-attention...\n", - "[INFO | DotProductAttention]: Running with FlashAttention backend\n", - "\n", - "Test passed.\n" - ] - } - ], - "source": [ - "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python example_attention.py" - ] - }, - { - "cell_type": "markdown", - "id": "8ca99461", - "metadata": {}, - "source": [ - "`NVTE_DEBUG_LEVEL=2` allows us to find out more about the backend selection logic. Users are encouraged to double check the `config` and provide it to the Transformer Engine team if they would like to file a bug. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d3637094", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Attention Backends\n", + "\n", + "Transformer Engine provides multiple attention backends for each supported framework. The framework-native backends provide a robust baseline, while the fused, GPU-optimized implementations offer more performance. For example, the flash-attention and cuDNN attention backends in PyTorch. The framework-native backends are often named with \"unfused\", while the more optimized backends are \"fused\" or \"flash\".\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
FrameworkBackend (Module Name)Module Location
PyTorchcuDNN attention (`FusedAttention`) [transformer_engine.pytorch.attention](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py)
flash-attention (`FlashAttention`)
\n", + " PyTorch-native attention (`UnfusedDotProductAttention`)\n", + "
JAXcuDNN attention (`_FusedDotProductAttention`)[transformer_engine.jax.flax.transformer](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/flax/transformer.py)
JAX-native attention (`_UnfusedDotProductAttention`)
" + ], + "id": "89a7d849" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Run cuDNN attention...\n", - "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", - "[DEBUG | DotProductAttention]: Disabling FlashAttention due to NVTE_FLASH_ATTN=0\n", - "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=False, FusedAttention=True (sub-backend 1), UnfusedDotProductAttention=True}\n", - "[DEBUG | DotProductAttention]: Selected backend = FusedAttention (sub-backend 1)\n", - "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", - "\n", - "Run flash-attention...\n", - "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", - "[DEBUG | DotProductAttention]: Disabling FusedAttention due to NVTE_FUSED_ATTN=0\n", - "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=True, FusedAttention=False, UnfusedDotProductAttention=True}\n", - "[DEBUG | DotProductAttention]: Selected backend = FlashAttention\n", - "[INFO | DotProductAttention]: Running with FlashAttention backend\n", - "\n", - "Test passed.\n" - ] - } - ], - "source": [ - "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py" - ] - }, - { - "cell_type": "markdown", - "id": "611d8fdb", - "metadata": {}, - "source": [ - "### 2.2 User Control\n", - "\n", - "Users usually do not need to worry about the backend selection. However, if there is a convergence or performance issue encountered, Transformer Engine provides a few other environment variables for users to experiment with different backends.\n", - "\n", - "**flash-attention or cuDNN attention:**\n", - "Users can enable/disable the flash-attention backend or cuDNN attention backend via the following two environment variables in PyTorch.\n", - "```\n", - "NVTE_FLASH_ATTN = 0 # disables flash-attention; default = 1\n", - "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n", - "```\n", - "\n", - "**cuDNN attention sub-backends:**\n", - "This environment variable allows users to express their preference of cuDNN attention sub-backends. However, the elected sub-backend will only be used *if* it is eligible, i.e. if it has support for the provided inputs and runtime environment.\n", - "```\n", - "NVTE_FUSED_ATTN_BACKEND = 1/2 # user preference of cuDNN sub-backend\n", - "```\n", - "\n", - "```\n", - "
\n", - "Note\n", - " \n", - "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, NVTE_FUSED_ATTN_BACKEND, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", - "
\n", - "\n", - "### 2.3 Example Tests\n", - "\n", - "Our [unit tests](https://github.com/NVIDIA/TransformerEngine/tree/main/tests) demonstrate the use of Transformer Engine dot product attention APIs. Users are encouraged to use them as a template when integrating Transformer Engine to their ML workflows.\n", - "\n", - "For example, in PyTorch, [test_dot_product_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) offers a variety of use cases of `pytorch.DotProductAttention`, from data types, model configs, checkpointing, to QKV layouts." - ] - }, - { - "cell_type": "markdown", - "id": "e60a2a3e", - "metadata": {}, - "source": [ - "## 3. Backend Support\n", - "\n", - "Transformer Engine supports commonly-used features such as self and cross attention, FP16/BF16 precisions, dropout, and checkpointing. But it also offers a range of other features. As of v2.0, Transformer Engine's attention backends have the following support matrix.\n", - "\n", - "| Attention Backend | Precision | Architecture | Sliding Window Attention | MQA/GQA | Multi-Latent Attention | Context Parallelism | Determinism Possible |\n", - "| :---------------- | :-------- | :----------- | :----------------------- | :------ | :--------------------- | :------------------ | :------------ |\n", - "| cuDNN attention (all frameworks) | BF16, FP16, FP8 (PyTorch only) | sm80+ | Yes (cuDNN 9.2+) | Yes | Yes | Yes (`bshd`,`sbhd`, `thd`) | Yes |\n", - "| flash-attention (PyTorch) | BF16, FP16 | sm80+ | Yes | Yes | Yes | Yes (`bshd`,`thd`) | Yes |\n", - "| Framework-native attention | BF16, FP16, FP32 | Any | No, unless used as a mask | Yes | Yes (PyTorch only) | No | Yes |\n", - "\n", - "Some unit tests are provided to serve as a starting point for integrating such features into users' models. For example,\n", - "- sliding window attention: [test_dpa_swa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", - "- MQA/GQA: [test_te_layer_mqa_gqa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", - "- Multi-Latent Attention: [test_dpa_mla](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", - "- context parallelism: [test_cp_with_fused_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py), [test_cp_with_flash_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py)" - ] - }, - { - "cell_type": "markdown", - "id": "fbdcb327", - "metadata": {}, - "source": [ - "### 3.1 QKV Layout\n", - "\n", - "Transformer Engine supports various layouts of the query `q`, key `k`, value `v` tensors. It has defined 15 QKV layouts, which are grouped into 3 QKV formats and 5 QKV layout groups to help with similar memory/computational operations across different layouts. The mapping relationships of these layouts and groups are,\n", - "\n", - "| `qkv_layout`         | `qkv_layout_group`=`3hd` | `h3d` | `hd_2hd` | `hd_h2d` | `hd_hd_hd` |\n", - "| ----------: | -----------: | -----: | ----------: | ----------: | -------------: |\n", - "| `qkv_format`=`sbhd` | `sb3hd` | `sbh3d` | `sbhd_sb2hd` | `sbhd_sbh2d` | `sbhd_sbhd_sbhd` |\n", - "| `bshd` | `bs3hd` | `bsh3d` | `bshd_bs2hd` | `bshd_bsh2d` | `bshd_bshd_bshd` |\n", - "| `thd` | `t3hd` | `th3d` | `thd_t2hd` | `thd_th2d` | `thd_thd_thd` |\n", - "\n", - "The notation system is that `b` stands for the batch size, `s` sequence length, `h` number of attention heads, `d` head dimension, and `t` the total number of tokens in the batch, i.e. `t = sum(s_i) for i in 0,...,b-1`. Here are a few examples of the layouts and their explanations to help clarify the definition.\n", - "\n", - "**qkv_layout=sb3hd:**\n", - "`q`, `k`, `v` are sequence first, i.e. `s` is the leading dimension in each tensor. They are different slices of one tensor `qkv`: `q, k, v = [qkv[:,:,i,:,:] for i in range(3)]`. They are interleaved at the `h * d` dimension.\n", - "\n", - "**qkv_layout=bshd_bsh2d:**\n", - "`q`, `k`, `v` are batch first, i.e. `b` is the leading dimension in each tensor. `q` is contiguous, and `k`, `v` are different slices of tensor `kv`: `k, v = [kv[:,:,:,i,:] for i in range(2)]`. `k`, `v` are interleaved at the `d` dimension.\n", - "\n", - "The `s` and `h` in `bsh2d` are the max sequence length and number of heads for `k`, `v`, which can be different from the `s` and `h` in `bshd` for `q`. We denoted them as the same for brevity reasons. Transformer Engine does differentiate their values for actual execution.\n", - "\n", - "**qkv_layout=thd_thd_thd:**\n", - "`q`, `k`, `v` have variable sequence lengths in a batch. They are all contiguous and have no interleaving.\n", - "\n", - "As of v2.0, Transformer Engine has the following support matrix.\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
BackendSupported QKV FormatsNotes
flash-attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
cuDNN attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
\n", - " JAX: `bs3hd`, `bshd_bs2hd`, `bshd_bshd_bshd` layouts\n", - "
Framework-native attention`bshd`, `sbhd`PyTorch, JAX: 2 formats, i.e. 10 layouts
\n", - "\n", - "Some example usage of the different layouts can be found at [test_dpa_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_dpa_qkv_layout_thd](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). Transformer Engine also provides a utility function [transformer_engine.pytorch.attention.dot_product_attention.utils.get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) to help determine which layout a set of `q`, `k`, `v` tensors have (PyTorch only).\n", - "\n", - "
\n", - "Note\n", - " \n", - "When RoPE is employed, the qkv_layout may change in Transformer Engine PyTorch through [get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py). This is due to the in-place nature of our RoPE implementations. We convert `q`, `k`, `v` tensors from their initial layout to the corresponding hd_hd_hd layout. For example, from sbh3d in pytorch.MultiHeadAttention before RoPE, to sbhd_sbhd_sbhd in pytorch.DotProductAttention after RoPE.\n", - "
\n" - ] - }, - { - "cell_type": "markdown", - "id": "855d9616", - "metadata": {}, - "source": [ - "### 3.2 Attention Mask\n", - "\n", - "Transformer Engine supports 7 mask types, and all the masks are defined as `True` masking out the corresponding element and `False` including the corresponding element in attention calculation.\n", - "\n", - "- `no_mask`, `padding`, `causal`, `causal_bottom_right`, `padding_causal`, `padding_causal_bottom_right`, `arbitrary`\n", - "\n", - "Different backends offer different support for attention mask. As of Transformer Engine 2.0,\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
BackendSupported Mask TypesRequires `attention_mask`
flash-attention
  • `no_mask`, `causal` (self-attention),
  • `padding`, `padding_causal` (self-attention),
  • `causal_bottom_right`, `padding_causal_bottom_right`
  • `no_mask`, `causal` `causal_bottom_right`: No
  • `padding`, `padding_causal`, `padding_causal_bottom_right`: Yes if `cu_seqlens` not provided
  • `arbitrary`: Yes
    • cuDNN attention
  • `no_mask`, `causal`,
  • `padding`, `padding_causal`,
  • `causal_bottom_right`, `padding_causal_bottom_right`
    • Framework-native attention
  • All (PyTorch)
  • `no_mask`, `causal`, `padding` (Jax)
    • \n", - "\n", - "**Padding masks:** For `padding`, `padding_causal`, `padding_causal_bottom_right` mask types, users need to provide sequence length information to help Transformer Engine figure out where each sequence ends in a batch. As of Transformer Engine 2.0, there are two options to do so in PyTorch and one in JAX.\n", - "\n", - "* PyTorch: When both options are provided by the user, `cu_seqlens` is preferred as there is no extra conversion needed.\n", - " - `cu_seqlens`: Users can provide cumulative sequence length tensors `cu_seqlens_q` and `cu_seqlens_kv` for `q` and `k`/`v` to the flash-attention or cuDNN attention backend. An example of `cu_seqlens` is `[0, 2, 6, 7]` for a batch of 3 `[aa000, bbbb0, c0000]`.\n", - " - `attention_mask`: Users can also provide `attention_mask` as an alternative, which will then be converted to `cu_seqlens`. For self-attention, `attention_mask` should be one single tensor of shape `[batch_size, 1, 1, seqlen_q]`, and for cross-attention, `attention_mask` should be a list of two tensors of shapes `[batch_size, 1, 1, seqlen_q]` and `[batch_size, 1, 1, seqlen_kv]`, respectively.\n", - "\n", - "\n", - "* JAX: Users should provide the `attention_mask` tensor of shape `[batch_size, 1, seqlen_q, seqlen_kv]`.\n", - "\n", - "**qkv_format=thd:** Transformer Engine extracts the max sequence length information from `q`, `k`, `v` if `max_seqlen_q` and `max_seqlen_kv` are not provided. This requires GPU-CPU copy and synchronization operations. For performance reasons, please set `max_seqlen_q` and `max_seqlen_kv` to their appropriate values for `thd` QKV format.\n", - "\n", - "**Arbitrary mask:** cuDNN does not support `Arbitrary` mask type as of v9.3. However, users can convert the mask to a regular `post_scale_bias` bias and achieve the same functionality. An example script for this conversion is [arbitrary_mask_to_post_scale_bias.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py).\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "a1f25a9b", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1.1 Flash vs. Non-Flash\n", + "\n", + "The attention calculation has quadratic computational and memory complexities to the sequence length. Its runtime and memory requirements quadruple, when the sequence length doubles. This presents a significant challenge to scale Transformer models up for longer contexts, in order to achieve higher model quality.\n", + "\n", + "Compared to the standard, non-flash algorithm, the flash algorithm [[2]](https://arxiv.org/abs/2205.14135) was proposed to reduce the memory scaling to linear and improve the computational efficiency through optimized memory accesses. It employs the following two distinctive techniques.\n", + "\n", + "- **Tiling:** The non-flash algorithm tries to process the query, key, value tensors in one single step, requiring large amounts of global memory and incurring high volumes of reads/writes between global memory and shared memory. The flash algorithm decomposes the input into several tiles, based on the available shared memory and register size, and it computes the softmax one tile at a time.\n", + "\n", + "- **Recomputation:** The non-flash algorithm stores the softmax matrix (quadratic to sequence length) to global memory for the backward pass, while the flash algorithm only saves the softmax normalization factors (linear to sequence length). This reduces the amount of memory required as well as the bandwidth utilization between global memory and shared memory. Even though there is extra computation incurred in order to recalculate the attention in the backward pass, the bandwidth savings still provide significant improvement in efficiency.\n", + "\n", + "
    \n", + "Note: \n", + " \n", + "Transformer Engine's flash-attention backend, available in PyTorch, and cuDNN attention backend (sub-backends 1 and 2), available in PyTorch and JAX, are both based on the flash algorithm.\n", + "
    \n" + ], + "id": "c90a2573" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "Run with post_scale_bias:\n", - "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", - "\n", - "Run with arbitrary mask:\n", - "[INFO | DotProductAttention]: Running with UnfusedDotProductAttention backend\n", - "\n", - "Test passed!\n" - ] + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1.2 flash-attention\n", + "\n", + "The flash-attention backend, available only in PyTorch, is a module wrapped around the public `flash-attn` package [[3]](https://github.com/Dao-AILab/flash-attention). \n", + "\n", + "The flash-attention backend supports `flash-attn`'s features as well as a few extra functionalities to facilitate the use of `flash-attn`, such as converting the `attention_mask` to cumulative sequence lengths `cu_seqlens` for `padding` mask use cases. Please see `transformer_engine.pytorch.attention.FlashAttention` for details.\n", + "\n", + "The `flash-attn` dependency is regularly updated in Transformer Engine. As of v2.0, Transformer Engine supports `flash-attn` 2.0.6+ (see [setup.py](https://github.com/NVIDIA/TransformerEngine/blob/main/setup.py)).\n", + "\n", + "To understand `flash-attn`'s performance, please refer to their benchmarks [here](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#performance).\n", + "\n", + "### 1.3 cuDNN Attention\n", + "\n", + "The cuDNN attention backend, available in PyTorch and JAX, offers another high-performance solution to the attention calculation. It requires [cuDNN](https://developer.nvidia.com/cudnn) to run, and has several sub-backends to support the different precisions and sequence lengths.\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    Sub-BackendAlgorithmPrecisionSequence LengthArchitectureAdditional info
    1FlashBF16/FP16 Any sm80+ [cuDNN](https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#fused-flash-attention-fprop),\n", + " [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention)\n", + "
    2FlashFP8 cuDNN pre-9.0: ≤512 cuDNN pre-9.0: sm90
    cuDNN 9.0+: Any cuDNN 9.0+: sm90+ cuDNN 9.0+: [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention-fp8)\n", + "
    \n", + "\n", + "The cuDNN attention backend and flash-attention backend have several notable differences. As of Transformer Engine 2.0, cuDNN 9.3 and `flash-attn` 2.4.2,\n", + "\n", + "- flash-attention only supports the PyTorch framework while cuDNN attention supports PyTorch and JAX.\n", + "- flash-attention supports BF16, FP16 precisions while cuDNN attention also supports FP8 (through its sub-backend 2).\n", + "- flash-attention supports `bshd`, `thd` input formats, without any transposes, and `sbhd` format, with transposes, while cuDNN attention supports all three formats without transposes (see Section 3.1 for more details).\n", + "- flash-attention does not support `post_scale_bias`, and cuDNN attention does.\n", + "- flash-attention supports KV-caching and paged attention, and cuDNN attention does not.\n", + "- flash-attention uses bottom right diagonal for `causal` mask in cross attention (see [change log](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#21-change-behavior-of-causal-flag)), and cuDNN attention supports both top left and bottom right.\n", + "- **Sliding window attention (SWA):** flash-attention has SWA(left, right) support for all mask types except top-left causal masks, with or without dropout, and without bias. cuDNN attention supports SWA(left, 0) starting from 9.2 and SWA(left, right) starting from 9.6, without dropout, and with `bias_type=\"no_bias\"`.\n", + "- flash-attention outperforms cuDNN attention on Ampere architectures, and cuDNN attention has 20-50% advantages on Hopper architectures, based on our benchmarks for a number of commonly-used model configurations.\n", + "\n", + "To compare cuDNN attention and flash-attention, users can modify the `model_configs` dictionary in [benchmarks/attention/benchmark_attention.py](https://github.com/NVIDIA/TransformerEngine/blob/main/benchmarks/attention/benchmark_attention.py) to collect performance numbers. The script runs each entry in `model_configs` for `num_iters` times, each time with one forward pass and one backward pass. Both backends are tried, and if one backend does not have support for the specific user input, the runtimes and speedups in the final table would be 0." + ], + "id": "b5ce567d" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_configs = {\n", + " # test: b, h, hg, d, sq, skv, p, mask, bias\n", + " \"test_0\": ModelConfig(2, 16, 16, 64, 512, 512, 0.0, \"no_mask\", \"no_bias\"), # short seq\n", + " \"test_1\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"no_bias\"), # longer seq, mask\n", + " \"test_2\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"post_scale_bias\"), # bias\n", + " \"test_3\": ModelConfig(2, 32, 4, 128, 8192, 8192, 0.0, \"causal\", \"no_bias\"), # GQA\n", + "}" + ], + "execution_count": null, + "outputs": [], + "id": "c5b8e3d7" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!cd ../../../benchmarks/attention/ && python benchmark_attention.py" + ], + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Device 0: NVIDIA H100 80GB HBM3 GPU, sm90 compute capability, 79.1GB memory\n", + "Running test_0 with cuDNN attention and flash-attention...\n", + "Running test_1 with cuDNN attention and flash-attention...\n", + "Running test_2 with cuDNN attention...\n", + "Running test_3 with cuDNN attention and flash-attention...\n", + "\n", + " cuDNN fwd+bwd (ms) flash-attn fwd+bwd (ms) cuDNN vs flash speedup\n", + "test_0 0.0340 0.0468 1.3786\n", + "test_1 0.3664 0.5850 1.5968\n", + "test_2 0.9332 0.0000 0.0000\n", + "test_3 7.4875 11.8879 1.5877\n" + ] + } + ], + "id": "50852cb5" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Backend Selection\n", + "\n", + "Given the various attention backends, Transformer Engine first determines which backends are eligible for the provided inputs and runtime environment, then applies a preference order among the eligible backends. Eligibility is affected by user environment variables, GPU architecture, installed `flash-attn` and cuDNN versions, data type and FP8 recipe, QKV layout, training or inference mode, dropout, and other attention features.\n", + "\n", + "In PyTorch, the candidates are FlashAttention (`flash-attn` v2, v3, or v4), FusedAttention (cuDNN sub-backends), and UnfusedDotProductAttention. Users can disable whole backend families with `NVTE_FLASH_ATTN`, `NVTE_FUSED_ATTN`, or `NVTE_UNFUSED_ATTN`. In JAX, Transformer Engine checks whether a cuDNN fused-attention kernel is available when `NVTE_FUSED_ATTN=1`; otherwise it falls back to the JAX-native implementation.\n", + "\n", + "At a high level, the architecture-specific PyTorch selection order is:\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    FrameworkSelection Order
    PyTorchsm8x (Ampere/Ada): flash-attention > cuDNN attention > PyTorch-native attention
    sm90 (Hopper): cuDNN attention > flash-attention > PyTorch-native attention
    sm100/sm120 (Blackwell): cuDNN attention > flash-attention > PyTorch-native attention
    cuDNN attention: BF16/FP16 uses sub-backend 1 when eligible; FP8 uses sub-backend 2 when enabled and eligible
    JAXcuDNN attention > JAX-native attention
    \n", + "\n", + "Within FlashAttention, TE uses the installed implementation that is supported for the architecture and input. FlashAttention 3 is Hopper-only (`sm90`). FlashAttention 4 supports `sm80`, `sm90`, `sm100`, and `sm120`; on Hopper, TE prefers FlashAttention 3 over FlashAttention 4 when both are installed and eligible. On Blackwell, FlashAttention 4 is the Blackwell-specific flash-attention path when installed and eligible, while FlashAttention 2 can still be eligible depending on the installed version and input configuration.\n", + "\n", + "Within cuDNN FusedAttention, TE asks the fused-attention helper which sub-backend is eligible. Sub-backend 1 is the BF16/FP16 flash-based path when available; sub-backend 2 is the FP8 path when FP8 DPA is enabled and the architecture, cuDNN version, and input configuration support it. Hopper supports eligible FP8 DPA through cuDNN sub-backend 2. In the current PyTorch selector, eligible FP8 DPA on Blackwell is an `sm100` path and is disabled on `sm120`.\n", + "\n", + "When all optimized backends are disabled or ineligible, TE falls back to UnfusedDotProductAttention if it is enabled. If no backend is eligible, backend selection returns no backend and the caller raises an error. As we monitor the performance of different backends, the selection logic may change." + ], + "id": "9a615119" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.1 Debug Information\n", + "\n", + "To find out which backend is being used during runtime, we have the following two debugging flags. Logging is done by using the `logging` package.\n", + "```\n", + "NVTE_DEBUG = 0/1 # disables/enables debugging\n", + "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n", + "```\n", + "
    \n", + "Note:\n", + " \n", + "These flags are supported in PyTorch only as of Transformer Engine 2.0. JAX support is expected to be added in the future.\n", + "
    " + ], + "id": "e6c0f3f0" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The example script [example_attention.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/example_attention.py) runs a very basic model with two attention backends, cuDNN attention and flash-attention. Here `NVTE_DEBUG_LEVEL=1` allows us to find out which backend/sub-backend is used in runtime." + ], + "id": "16660323" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python example_attention.py" + ], + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "text": [ + "\n", + "Run cuDNN attention...\n", + "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", + "\n", + "Run flash-attention...\n", + "[INFO | DotProductAttention]: Running with FlashAttention backend\n", + "\n", + "Test passed.\n" + ] + } + ], + "id": "906b8cf1" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`NVTE_DEBUG_LEVEL=2` allows us to find out more about the backend selection logic. Users are encouraged to double check the `config` and provide it to the Transformer Engine team if they would like to file a bug. " + ], + "id": "8ca99461" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py" + ], + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "text": [ + "\n", + "Run cuDNN attention...\n", + "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", + "[DEBUG | DotProductAttention]: Disabling FlashAttention due to NVTE_FLASH_ATTN=0\n", + "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=False, FusedAttention=True (sub-backend 1), UnfusedDotProductAttention=True}\n", + "[DEBUG | DotProductAttention]: Selected backend = FusedAttention (sub-backend 1)\n", + "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", + "\n", + "Run flash-attention...\n", + "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", + "[DEBUG | DotProductAttention]: Disabling FusedAttention due to NVTE_FUSED_ATTN=0\n", + "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=True, FusedAttention=False, UnfusedDotProductAttention=True}\n", + "[DEBUG | DotProductAttention]: Selected backend = FlashAttention\n", + "[INFO | DotProductAttention]: Running with FlashAttention backend\n", + "\n", + "Test passed.\n" + ] + } + ], + "id": "d3637094" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.2 User Control\n", + "\n", + "Users usually do not need to worry about the backend selection. However, if there is a convergence or performance issue encountered, Transformer Engine provides a few other environment variables for users to experiment with different backends.\n", + "\n", + "**flash-attention or cuDNN attention:**\n", + "Users can enable/disable the flash-attention backend or cuDNN attention backend via the following two environment variables in PyTorch.\n", + "```\n", + "NVTE_FLASH_ATTN = 0 # disables flash-attention; default = 1\n", + "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n", + "```\n", + "\n", + "```\n", + "
    \n", + "Note\n", + " \n", + "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", + "
    \n", + "\n", + "### 2.3 Example Tests\n", + "\n", + "Our [unit tests](https://github.com/NVIDIA/TransformerEngine/tree/main/tests) demonstrate the use of Transformer Engine dot product attention APIs. Users are encouraged to use them as a template when integrating Transformer Engine to their ML workflows.\n", + "\n", + "For example, in PyTorch, [test_dot_product_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) offers a variety of use cases of `pytorch.DotProductAttention`, from data types, model configs, checkpointing, to QKV layouts." + ], + "id": "611d8fdb" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Backend Support\n", + "\n", + "Transformer Engine supports commonly-used features such as self and cross attention, FP16/BF16 precisions, dropout, and checkpointing. But it also offers a range of other features. As of v2.0, Transformer Engine's attention backends have the following support matrix.\n", + "\n", + "| Attention Backend | Precision | Architecture | Sliding Window Attention | MQA/GQA | Multi-Latent Attention | Context Parallelism | Determinism Possible |\n", + "| :---------------- | :-------- | :----------- | :----------------------- | :------ | :--------------------- | :------------------ | :------------ |\n", + "| cuDNN attention (all frameworks) | BF16, FP16, FP8 (PyTorch only) | sm80+ | Yes (cuDNN 9.2+) | Yes | Yes | Yes (`bshd`,`sbhd`, `thd`) | Yes |\n", + "| flash-attention (PyTorch) | BF16, FP16 | sm80+ | Yes | Yes | Yes | Yes (`bshd`,`thd`) | Yes |\n", + "| Framework-native attention | BF16, FP16, FP32 | Any | No, unless used as a mask | Yes | Yes (PyTorch only) | No | Yes |\n", + "\n", + "Some unit tests are provided to serve as a starting point for integrating such features into users' models. For example,\n", + "- sliding window attention: [test_dpa_swa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", + "- MQA/GQA: [test_te_layer_mqa_gqa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", + "- Multi-Latent Attention: [test_dpa_mla](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", + "- context parallelism: [test_cp_with_fused_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py), [test_cp_with_flash_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py)" + ], + "id": "e60a2a3e" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.1 QKV Layout\n", + "\n", + "Transformer Engine supports various layouts of the query `q`, key `k`, value `v` tensors. It has defined 15 QKV layouts, which are grouped into 3 QKV formats and 5 QKV layout groups to help with similar memory/computational operations across different layouts. The mapping relationships of these layouts and groups are,\n", + "\n", + "| `qkv_layout`         | `qkv_layout_group`=`3hd` | `h3d` | `hd_2hd` | `hd_h2d` | `hd_hd_hd` |\n", + "| ----------: | -----------: | -----: | ----------: | ----------: | -------------: |\n", + "| `qkv_format`=`sbhd` | `sb3hd` | `sbh3d` | `sbhd_sb2hd` | `sbhd_sbh2d` | `sbhd_sbhd_sbhd` |\n", + "| `bshd` | `bs3hd` | `bsh3d` | `bshd_bs2hd` | `bshd_bsh2d` | `bshd_bshd_bshd` |\n", + "| `thd` | `t3hd` | `th3d` | `thd_t2hd` | `thd_th2d` | `thd_thd_thd` |\n", + "\n", + "The notation system is that `b` stands for the batch size, `s` sequence length, `h` number of attention heads, `d` head dimension, and `t` the total number of tokens in the batch, i.e. `t = sum(s_i) for i in 0,...,b-1`. Here are a few examples of the layouts and their explanations to help clarify the definition.\n", + "\n", + "**qkv_layout=sb3hd:**\n", + "`q`, `k`, `v` are sequence first, i.e. `s` is the leading dimension in each tensor. They are different slices of one tensor `qkv`: `q, k, v = [qkv[:,:,i,:,:] for i in range(3)]`. They are interleaved at the `h * d` dimension.\n", + "\n", + "**qkv_layout=bshd_bsh2d:**\n", + "`q`, `k`, `v` are batch first, i.e. `b` is the leading dimension in each tensor. `q` is contiguous, and `k`, `v` are different slices of tensor `kv`: `k, v = [kv[:,:,:,i,:] for i in range(2)]`. `k`, `v` are interleaved at the `d` dimension.\n", + "\n", + "The `s` and `h` in `bsh2d` are the max sequence length and number of heads for `k`, `v`, which can be different from the `s` and `h` in `bshd` for `q`. We denoted them as the same for brevity reasons. Transformer Engine does differentiate their values for actual execution.\n", + "\n", + "**qkv_layout=thd_thd_thd:**\n", + "`q`, `k`, `v` have variable sequence lengths in a batch. They are all contiguous and have no interleaving.\n", + "\n", + "As of v2.0, Transformer Engine has the following support matrix.\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    BackendSupported QKV FormatsNotes
    flash-attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
    cuDNN attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
    \n", + " JAX: `bs3hd`, `bshd_bs2hd`, `bshd_bshd_bshd` layouts\n", + "
    Framework-native attention`bshd`, `sbhd`PyTorch, JAX: 2 formats, i.e. 10 layouts
    \n", + "\n", + "Some example usage of the different layouts can be found at [test_dpa_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_dpa_qkv_layout_thd](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). Transformer Engine also provides a utility function [transformer_engine.pytorch.attention.dot_product_attention.utils.get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) to help determine which layout a set of `q`, `k`, `v` tensors have (PyTorch only).\n", + "\n", + "
    \n", + "Note\n", + " \n", + "When RoPE is employed, the qkv_layout may change in Transformer Engine PyTorch through [get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py). This is due to the in-place nature of our RoPE implementations. We convert `q`, `k`, `v` tensors from their initial layout to the corresponding hd_hd_hd layout. For example, from sbh3d in pytorch.MultiHeadAttention before RoPE, to sbhd_sbhd_sbhd in pytorch.DotProductAttention after RoPE.\n", + "
    \n" + ], + "id": "fbdcb327" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.2 Attention Mask\n", + "\n", + "Transformer Engine supports 7 mask types, and all the masks are defined as `True` masking out the corresponding element and `False` including the corresponding element in attention calculation.\n", + "\n", + "- `no_mask`, `padding`, `causal`, `causal_bottom_right`, `padding_causal`, `padding_causal_bottom_right`, `arbitrary`\n", + "\n", + "Different backends offer different support for attention mask. As of Transformer Engine 2.0,\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    BackendSupported Mask TypesRequires `attention_mask`
    flash-attention
  • `no_mask`, `causal` (self-attention),
  • `padding`, `padding_causal` (self-attention),
  • `causal_bottom_right`, `padding_causal_bottom_right`
  • `no_mask`, `causal` `causal_bottom_right`: No
  • `padding`, `padding_causal`, `padding_causal_bottom_right`: Yes if `cu_seqlens` not provided
  • `arbitrary`: Yes
    • cuDNN attention
  • `no_mask`, `causal`,
  • `padding`, `padding_causal`,
  • `causal_bottom_right`, `padding_causal_bottom_right`
    • Framework-native attention
  • All (PyTorch)
  • `no_mask`, `causal`, `padding` (Jax)
    • \n", + "\n", + "**Padding masks:** For `padding`, `padding_causal`, `padding_causal_bottom_right` mask types, users need to provide sequence length information to help Transformer Engine figure out where each sequence ends in a batch. As of Transformer Engine 2.0, there are two options to do so in PyTorch and one in JAX.\n", + "\n", + "* PyTorch: When both options are provided by the user, `cu_seqlens` is preferred as there is no extra conversion needed.\n", + " - `cu_seqlens`: Users can provide cumulative sequence length tensors `cu_seqlens_q` and `cu_seqlens_kv` for `q` and `k`/`v` to the flash-attention or cuDNN attention backend. An example of `cu_seqlens` is `[0, 2, 6, 7]` for a batch of 3 `[aa000, bbbb0, c0000]`.\n", + " - `attention_mask`: Users can also provide `attention_mask` as an alternative, which will then be converted to `cu_seqlens`. For self-attention, `attention_mask` should be one single tensor of shape `[batch_size, 1, 1, seqlen_q]`, and for cross-attention, `attention_mask` should be a list of two tensors of shapes `[batch_size, 1, 1, seqlen_q]` and `[batch_size, 1, 1, seqlen_kv]`, respectively.\n", + "\n", + "\n", + "* JAX: Users should provide the `attention_mask` tensor of shape `[batch_size, 1, seqlen_q, seqlen_kv]`.\n", + "\n", + "**qkv_format=thd:** Transformer Engine extracts the max sequence length information from `q`, `k`, `v` if `max_seqlen_q` and `max_seqlen_kv` are not provided. This requires GPU-CPU copy and synchronization operations. For performance reasons, please set `max_seqlen_q` and `max_seqlen_kv` to their appropriate values for `thd` QKV format.\n", + "\n", + "**Arbitrary mask:** cuDNN does not support `Arbitrary` mask type as of v9.3. However, users can convert the mask to a regular `post_scale_bias` bias and achieve the same functionality. An example script for this conversion is [arbitrary_mask_to_post_scale_bias.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py).\n" + ], + "id": "855d9616" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python arbitrary_mask_to_post_scale_bias.py" + ], + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Run with post_scale_bias:\n", + "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", + "\n", + "Run with arbitrary mask:\n", + "[INFO | DotProductAttention]: Running with UnfusedDotProductAttention backend\n", + "\n", + "Test passed!\n" + ] + } + ], + "id": "a1f25a9b" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Some more examples of running Transformer Engine with different attention masks can be found at [test_dpa_mask](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py).\n", + "\n", + "### 3.3 Attention Bias\n", + "\n", + "Transformer Engine supports 4 attention bias types, `no_bias`, `pre_scale_bias`, `post_scale_bias`, and `ALiBi` (with/without custom slopes). As of Transformer Engine 2.0, their support matrix is as follows.\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    BackendBias TypeBias ShapeBias Data TypeArchitecture
    flash-attention`no_bias`, `ALiBi` (with slopes)N/AALiBi slopes: FP32sm80+
    cuDNN attentionPyTorch: `no_bias`, `post_scale_bias`, `ALiBi` (without slopes)`post_scale_bias`: BHSS, 1HSS, B1SS, 11SS for forward, 1HSS for backward`post_scale_bias`: same as QKV typecuDNN 8.9.6+: sm90
    JAX: `no_bias`, `post_scale_bias`ALiBi slopes: FP32cuDNN 9.0+: sm80+
    Framework-native attention`no_bias`, `pre_scale_bias`, `post_scale_bias``post_scale_bias`: BHSS, 1HSS, B1SS, 11SS `post_scale_bias`: same as QKV typesm80+
    \n", + "\n", + "The flash-attention backend enables `ALiBi` by asking user to pass in an `alibi_slopes` tensor, which can be the default slopes of vanilla ALiBi, or user-defined slopes. On the other hand, cuDNN attention supports `ALiBi` by taking in a `Boolean` flag, and it only supports vanilla ALiBi as of cuDNN 9.0.\n", + "\n", + "The framework-native backends do not explicitly support `ALiBi`, but users can convert `ALiBi` to a regular `post_scale_bias` bias to achieve the same effect. In PyTorch, this utility function, `transformer_engine.pytorch.attention.get_alibi`, can be used to help with the conversion.\n", + "\n", + "More examples of how to use the various attention biases are at [test_dpa_bias](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)." + ], + "id": "dda4a589" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.4 FP8 Attention\n", + "\n", + "A unique feature of Transformer Engine is its FP8 support, not only for the `Linear` layers but also for dot product attention. Transformer Engine's FP8 attention support is through its cuDNN attention sub-backend 2. Recall Figure 1: the two `MatMul` operations are performed in FP8 for computational efficiency, and the `SoftMax` operation is performed in FP32 for numerical accuracy.\n", + "\n", + "Transformer Engine supports FP8 attention through its [C APIs](../../api/c/fused_attn.rst), and [PyTorch API](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention), as of v2.0. Its PyTorch API offers two options, both controlled through the FP8 recipe definition, `transformer_engine.common.recipe.DelayedScaling`.\n", + "\n", + "- `DelayedScaling.fp8_dpa=True (default=False)`: This enables the use of cuDNN attention sub-backend 2, when it does support the provided user inputs. The `FusedAttention` module for cuDNN attention takes FP16 or BF16 tensors as inputs, performs dot product attention in FP8, and returns attention logits in FP16 or BF16 (same as the input type). Casting operations are required to cast tensors to FP8 at the beginning, and back to FP16/BF16 at the end of the module.\n", + "\n", + "- `DelayedScaling.fp8_mha=True (default=False)`: This option, on top of `fp8_dpa=True`, removes the casting operations at the beginning and end of the `FusedAttention` module. This feature is experimental. \n", + "\n", + "Examples of using the two features are available at [test_dpa_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_mha_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). To disable FP8 attention for backward and only use it for forward, users can also set `NVTE_FP8_DPA_BWD=0 (default=1)`." + ], + "id": "a0702339" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" } - ], - "source": [ - "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python arbitrary_mask_to_post_scale_bias.py" - ] - }, - { - "cell_type": "markdown", - "id": "dda4a589", - "metadata": {}, - "source": [ - "Some more examples of running Transformer Engine with different attention masks can be found at [test_dpa_mask](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py).\n", - "\n", - "### 3.3 Attention Bias\n", - "\n", - "Transformer Engine supports 4 attention bias types, `no_bias`, `pre_scale_bias`, `post_scale_bias`, and `ALiBi` (with/without custom slopes). As of Transformer Engine 2.0, their support matrix is as follows.\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    BackendBias TypeBias ShapeBias Data TypeArchitecture
    flash-attention`no_bias`, `ALiBi` (with slopes)N/AALiBi slopes: FP32sm80+
    cuDNN attentionPyTorch: `no_bias`, `post_scale_bias`, `ALiBi` (without slopes)`post_scale_bias`: BHSS, 1HSS, B1SS, 11SS for forward, 1HSS for backward`post_scale_bias`: same as QKV typecuDNN 8.9.6+: sm90
    JAX: `no_bias`, `post_scale_bias`ALiBi slopes: FP32cuDNN 9.0+: sm80+
    Framework-native attention`no_bias`, `pre_scale_bias`, `post_scale_bias``post_scale_bias`: BHSS, 1HSS, B1SS, 11SS `post_scale_bias`: same as QKV typesm80+
    \n", - "\n", - "The flash-attention backend enables `ALiBi` by asking user to pass in an `alibi_slopes` tensor, which can be the default slopes of vanilla ALiBi, or user-defined slopes. On the other hand, cuDNN attention supports `ALiBi` by taking in a `Boolean` flag, and it only supports vanilla ALiBi as of cuDNN 9.0.\n", - "\n", - "The framework-native backends do not explicitly support `ALiBi`, but users can convert `ALiBi` to a regular `post_scale_bias` bias to achieve the same effect. In PyTorch, this utility function, `transformer_engine.pytorch.attention.get_alibi`, can be used to help with the conversion.\n", - "\n", - "More examples of how to use the various attention biases are at [test_dpa_bias](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)." - ] - }, - { - "cell_type": "markdown", - "id": "a0702339", - "metadata": {}, - "source": [ - "### 3.4 FP8 Attention\n", - "\n", - "A unique feature of Transformer Engine is its FP8 support, not only for the `Linear` layers but also for dot product attention. Transformer Engine's FP8 attention support is through its cuDNN attention sub-backend 2. Recall Figure 1: the two `MatMul` operations are performed in FP8 for computational efficiency, and the `SoftMax` operation is performed in FP32 for numerical accuracy.\n", - "\n", - "Transformer Engine supports FP8 attention through its [C APIs](../../api/c/fused_attn.rst), and [PyTorch API](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention), as of v2.0. Its PyTorch API offers two options, both controlled through the FP8 recipe definition, `transformer_engine.common.recipe.DelayedScaling`.\n", - "\n", - "- `DelayedScaling.fp8_dpa=True (default=False)`: This enables the use of cuDNN attention sub-backend 2, when it does support the provided user inputs. The `FusedAttention` module for cuDNN attention takes FP16 or BF16 tensors as inputs, performs dot product attention in FP8, and returns attention logits in FP16 or BF16 (same as the input type). Casting operations are required to cast tensors to FP8 at the beginning, and back to FP16/BF16 at the end of the module.\n", - "\n", - "- `DelayedScaling.fp8_mha=True (default=False)`: This option, on top of `fp8_dpa=True`, removes the casting operations at the beginning and end of the `FusedAttention` module. This feature is experimental. \n", - "\n", - "Examples of using the two features are available at [test_dpa_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_mha_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). To disable FP8 attention for backward and only use it for forward, users can also set `NVTE_FP8_DPA_BWD=0 (default=1)`." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/docs/installation.rst b/docs/installation.rst index 0271af7fcc..cc48a0adac 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -14,7 +14,7 @@ Prerequisites 1. Linux x86_64 2. `CUDA 12.1+ (12.8+ for Blackwell support) `__ 3. |driver link|_ supporting CUDA 12.1 or later. -4. `cuDNN 9.11 `__ or later. +4. `cuDNN 9.3 `__ or later. If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within `CUDA_HOME`, set diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index e63b7b7b04..731958ec43 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -23,9 +23,6 @@ is_fp8_available, is_bf16_available, ) -from transformer_engine.pytorch.attention.dot_product_attention import ( - _attention_backends, -) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, check_set_window_size, @@ -1028,8 +1025,6 @@ def _run_dot_product_attention( os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - # Create seqlens qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) if "padding" in config.attn_mask_type or qkv_format == "thd": @@ -1584,8 +1579,6 @@ def _run_transformer_layer( os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - # Create input tensor if qkv_format == "sbhd": inp = torch.randn( @@ -2040,7 +2033,6 @@ def test_mha_fp8_vs_f16( os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") flash_attn_fwd_fp8, param_names, flash_attn_bwd_fp8 = _run_mha_fp8_vs_f16( dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe @@ -2050,7 +2042,6 @@ def test_mha_fp8_vs_f16( os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") fused_attn_fwd_fp8, param_names, fused_attn_bwd_fp8 = _run_mha_fp8_vs_f16( dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe @@ -2060,7 +2051,6 @@ def test_mha_fp8_vs_f16( os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = False") fused_attn_fwd_f16, param_names, fused_attn_bwd_f16 = _run_mha_fp8_vs_f16( dtype, config, False, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe @@ -2300,7 +2290,6 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FlashAttention)") flash_attn_fwd_fp8, flash_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( dtype, config, True, qkv_layout, is_training, fp8_recipe @@ -2310,7 +2299,6 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (UnfusedDotProductAttention)") unfused_attn_fwd_fp8, unfused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( dtype, config, True, qkv_layout, is_training, fp8_recipe @@ -2320,7 +2308,6 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FusedAttention)") fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( dtype, config, True, qkv_layout, is_training, fp8_recipe @@ -2650,8 +2637,6 @@ def _run_custom_mha_fp8(dtype, config, backend): os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - inp = 0.0001 * torch.randint( -100, 100, @@ -2708,8 +2693,6 @@ def _run_ref_mha_f16(dtype, config, backend): os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - inp = torch.load("qkv.pt").to(device="cuda") inp.requires_grad = True seqlens = torch.full([config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda") diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 34bfe7b939..fe8f416af4 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -453,12 +453,14 @@ def test(): if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() - for i in backends: - os.environ["NVTE_FUSED_ATTN_BACKEND"] = str(i) - _attention_backends["backend_selection_requires_update"] = True - available_backends, flash_attention_backend, fused_attention_backend = test() - if fused_attention_backend == FusedAttnBackend[backends[i]]: - fused_attn_backends.append(fused_attention_backend) + # F16_arbitrary_seqlen and FP8 are mutually exclusive for a given config (selected by the + # FP8/dtype gate), so a single probe returns the one applicable fused sub-backend. The old + # loop force-set NVTE_FUSED_ATTN_BACKEND per sub-backend, but the refactored backend selection + # no longer reads that env var, so the forcing was inert (and leaked the env var). + _attention_backends["backend_selection_requires_update"] = True + available_backends, flash_attention_backend, fused_attention_backend = test() + if fused_attention_backend in (FusedAttnBackend[name] for name in backends.values()): + fused_attn_backends.append(fused_attention_backend) return available_backends, flash_attention_backend, fused_attn_backends diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index ba86420ef5..57cdd385b2 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -13,17 +13,8 @@ if (CMAKE_BUILD_TYPE STREQUAL "Debug") endif() # Hide non-necessary symbols in shared object. -# [GRAPH-DEBUG] -DNVTE_GRAPH_DEBUG_SYMBOLS=ON keeps + exports internal symbols so -# backtrace_symbols() in fused_attn/graph_debug.h can name fused-attn frames. Remove after -# verification (revert to the two unconditional --version-script lines). -option(NVTE_GRAPH_DEBUG_SYMBOLS "Export all symbols for readable backtraces" OFF) -if (NOT NVTE_GRAPH_DEBUG_SYMBOLS) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") - set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") -else() - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -rdynamic -Wl,--export-dynamic") - set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -rdynamic -Wl,--export-dynamic") -endif() +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") +set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/libtransformer_engine.version") # Transformer Engine library project(transformer_engine LANGUAGES CUDA CXX) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 16727f80e8..cf5e0f3475 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -29,11 +29,11 @@ void uint8_to_bool(const void *in, bool &out) { 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 FusedAttnConfig::derive() { const int64_t b = static_cast(batch_size); @@ -332,20 +332,22 @@ FusedAttnConfig FusedAttnBwdParams::make_config() const { return cfg; } +} // namespace fused_attn } // namespace transformer_engine NVTEFusedAttnConfig nvte_create_fused_attn_config() { - return new transformer_engine::FusedAttnConfig{}; + return new transformer_engine::fused_attn::FusedAttnConfig{}; } void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config) { - delete transformer_engine::get_fused_attn_config_mutable(config); + delete transformer_engine::fused_attn::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; + using namespace transformer_engine::fused_attn; NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, "Invalid NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); @@ -498,6 +500,7 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, NVTEFusedAttnConfigAttribute attr, const void *buf, size_t size_in_bytes) { using namespace transformer_engine; + using namespace transformer_engine::fused_attn; NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, "Invalid NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); @@ -642,17 +645,18 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, } NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params() { - return new transformer_engine::FusedAttnFwdParams{}; + return new transformer_engine::fused_attn::FusedAttnFwdParams{}; } void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { - delete transformer_engine::get_fused_attn_fwd_params_mutable(params); + delete transformer_engine::fused_attn::get_fused_attn_fwd_params_mutable(params); } void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, NVTEFusedAttnFwdParamsAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written) { using namespace transformer_engine; + using namespace transformer_engine::fused_attn; NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; @@ -774,6 +778,7 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, NVTEFusedAttnFwdParamsAttribute attr, const void *buf, size_t size_in_bytes) { using namespace transformer_engine; + using namespace transformer_engine::fused_attn; NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; @@ -887,17 +892,18 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, } NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params() { - return new transformer_engine::FusedAttnBwdParams{}; + return new transformer_engine::fused_attn::FusedAttnBwdParams{}; } void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { - delete transformer_engine::get_fused_attn_bwd_params_mutable(params); + delete transformer_engine::fused_attn::get_fused_attn_bwd_params_mutable(params); } void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written) { using namespace transformer_engine; + using namespace transformer_engine::fused_attn; NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; @@ -1031,6 +1037,7 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, const void *buf, size_t size_in_bytes) { using namespace transformer_engine; + using namespace transformer_engine::fused_attn; NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 4839a73af7..ab01de9b91 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -17,6 +17,7 @@ #include "transformer_engine/fused_attn.h" namespace transformer_engine { +namespace fused_attn { struct FusedAttnConfig { // basic attention settings @@ -374,6 +375,7 @@ inline FusedAttnBwdParams *get_fused_attn_bwd_params_mutable(NVTEFusedAttnBwdPar return reinterpret_cast(params); } +} // namespace fused_attn } // 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 c6afbe30f1..159063fd27 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -247,6 +247,7 @@ void set_message(const char **message, std::string reason) { NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, const char **message) { using namespace transformer_engine; + using namespace transformer_engine::fused_attn; const FusedAttnConfig &cfg = *get_fused_attn_config(config); set_message(message, ""); @@ -348,7 +349,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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) { - transformer_engine::FusedAttnConfig cfg{}; + transformer_engine::fused_attn::FusedAttnConfig cfg{}; cfg.qkv_layout = qkv_layout; cfg.bias_type = bias_type; cfg.attn_mask_type = attn_mask_type; @@ -378,6 +379,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; + using namespace transformer_engine::fused_attn; const FusedAttnFwdParams &p = *get_fused_attn_fwd_params(params); const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); @@ -432,7 +434,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd); - transformer_engine::FusedAttnFwdParams p{}; + transformer_engine::fused_attn::FusedAttnFwdParams p{}; p.Q = Q; p.K = K; p.V = V; @@ -472,6 +474,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { NVTE_API_CALL(nvte_fused_attn_bwd_v2); using namespace transformer_engine; + using namespace transformer_engine::fused_attn; const FusedAttnBwdParams &p = *get_fused_attn_bwd_params(params); const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); @@ -551,7 +554,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_bwd); - transformer_engine::FusedAttnBwdParams p{}; + transformer_engine::fused_attn::FusedAttnBwdParams p{}; p.Q = Q; p.K = K; p.V = V; 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 460cd234b2..b0316b83ff 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 @@ -18,35 +18,9 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" -#include "graph_debug.h" // [GRAPH-DEBUG] +#include "graph_cache_debug.h" // [FUSED-ATTN-CACHE] #include "utils.h" -#define Q_ID 1 -#define K_ID 2 -#define V_ID 3 -#define O_ID 4 -#define S_ID 5 -#define B_ID 6 -#define D_CONST_ID 7 -#define S_CONST_ID 8 -#define Q_SEQLEN_ID 9 -#define K_SEQLEN_ID 10 -#define dQ_ID 11 -#define dK_ID 12 -#define dV_ID 13 -#define dO_ID 14 -#define MASK_VAL_ID 15 -#define dS_ID 16 -#define D_SEED_ID 17 -#define D_OFFSET_ID 18 -#define S_STATS_ID 19 -#define S_SUM_ID 20 -#define SCALE_PROB 21 -#define K_TRANSPOSE_ID 22 -#define dQ_ACCUM_ID 23 - -#define VIRTUAL_ID 30 - namespace transformer_engine { namespace fused_attn { @@ -87,7 +61,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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; @@ -103,8 +76,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( bool is_padding = cfg.is_padding; bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (is_training && dropout_probability != 0.0f); - NVTE_QKV_Format q_format = cfg.q_format; - NVTE_QKV_Format kv_format = cfg.kv_format; bool is_ragged_q = cfg.is_ragged_q; bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); @@ -186,12 +157,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // [SHARED-CACHE] Process-wide graph cache (was `static thread_local`) so a compiled graph // is reused across threads instead of rebuilt per thread. Safe because cuDNN >= 9.0 allows // concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe - // execute(); the static_asserts below fail the build loudly on an older toolkit. - static_assert(CUDNN_VERSION >= 91100, - "[SHARED-CACHE] shared fused-attn graph cache requires cuDNN >= 9.11 " - "(TE minimum supported cuDNN version)"); - static_assert(CUDNN_FRONTEND_VERSION >= 12500, - "[SHARED-CACHE] shared fused-attn graph cache requires cudnn-frontend >= 1.25.0"); + // execute(). The TE minimum-cuDNN-version bump that formalizes this requirement is a follow-up PR. static CacheType sdpa_f16_fprop_cache; static std::mutex sdpa_f16_fprop_cache_mutex; @@ -208,14 +174,14 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] - sm_arch_ != 120) { // [GRAPH-DEBUG] - fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] - "fwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), - /*legacy=*/!use_cu_seqlens_directly); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + graph_cache_debug::note_cache_lookup("fwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [FUSED-ATTN-CACHE] + sm_arch_ != 120) { // [FUSED-ATTN-CACHE] + graph_cache_debug::note_thd_lookup( // [FUSED-ATTN-CACHE] + "fwd", cache_hit, !cache_hit || graph_cache_debug::cache_disabled(), + /*legacy=*/!use_cu_seqlens_directly); // [FUSED-ATTN-CACHE] + } // [FUSED-ATTN-CACHE] + if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] return cached_graph; } @@ -482,31 +448,23 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] - GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); - GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); - GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) - GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->validate()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - fused_attn_graph_debug::note_fwd_build(); // [GRAPH-DEBUG] - if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] - std::vector serialized_graph; // [GRAPH-DEBUG] - if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] - fused_attn_graph_debug::note_graph_size("fwd", serialized_graph.size()); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] + graph_cache_debug::note_fwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - fused_attn_graph_debug::note_cache_size("fwd", cache.size()); // [GRAPH-DEBUG] - return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; } }; @@ -541,7 +499,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - fused_attn_graph_debug::note_fwd_exec(); // [GRAPH-DEBUG] + graph_cache_debug::note_fwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -686,9 +644,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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; @@ -705,8 +660,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( bool is_padding = cfg.is_padding; bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (dropout_probability != 0.0f); - NVTE_QKV_Format q_format = cfg.q_format; - NVTE_QKV_Format kv_format = cfg.kv_format; bool is_ragged_q = cfg.is_ragged_q; bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); @@ -777,15 +730,15 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [GRAPH-DEBUG] - sm_arch_ != 120) { // [GRAPH-DEBUG] + graph_cache_debug::note_cache_lookup("bwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [FUSED-ATTN-CACHE] + sm_arch_ != 120) { // [FUSED-ATTN-CACHE] // The backward impl has no cu_seqlens-directly path; it always buckets the batch. - fused_attn_graph_debug::note_thd_lookup( // [GRAPH-DEBUG] - "bwd", cache_hit, !cache_hit || fused_attn_graph_debug::cache_disabled(), - /*legacy=*/true); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + graph_cache_debug::note_thd_lookup( // [FUSED-ATTN-CACHE] + "bwd", cache_hit, !cache_hit || graph_cache_debug::cache_disabled(), + /*legacy=*/true); // [FUSED-ATTN-CACHE] + } // [FUSED-ATTN-CACHE] + if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] return cached_graph; } @@ -1024,30 +977,22 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] - GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); - GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); - GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) - GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->validate()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - fused_attn_graph_debug::note_bwd_build(); // [GRAPH-DEBUG] - if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] - std::vector serialized_graph; // [GRAPH-DEBUG] - if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] - fused_attn_graph_debug::note_graph_size("bwd", serialized_graph.size()); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] + graph_cache_debug::note_bwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - fused_attn_graph_debug::note_cache_size("bwd", cache.size()); // [GRAPH-DEBUG] - return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; } }; @@ -1077,7 +1022,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - fused_attn_graph_debug::note_bwd_exec(); // [GRAPH-DEBUG] + graph_cache_debug::note_bwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1193,11 +1138,7 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i 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; 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 3f22f131d8..c493b5ee1b 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 @@ -8,8 +8,8 @@ * \brief Functions for fused attention with seqlen > 512 */ -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ #include @@ -20,7 +20,7 @@ #include "transformer_engine/fused_attn.h" namespace transformer_engine { -void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *input_Q, +void fused_attn_arbitrary_seqlen_fwd(const fused_attn::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, @@ -30,7 +30,7 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i const Tensor *page_table_v, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); -void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *input_Q, +void fused_attn_arbitrary_seqlen_bwd(const fused_attn::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, @@ -44,13 +44,13 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i // 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 explaining why it is not supported. -std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); +std::string is_supported_f16_fwd(const fused_attn::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 explaining why it is not supported. -std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); +std::string is_supported_f16_bwd(const fused_attn::FusedAttnConfig &cfg, cudnnHandle_t handle); } // namespace transformer_engine -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_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 5eddd8fa75..a98a2f1950 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -5,13 +5,13 @@ ************************************************************************/ #include // [SHARED-CACHE] -#include // [GRAPH-DEBUG] serialized-size probe +#include // [FUSED-ATTN-CACHE] serialized-size probe #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" -#include "graph_debug.h" // [GRAPH-DEBUG] +#include "graph_cache_debug.h" // [FUSED-ATTN-CACHE] #include "utils.h" namespace transformer_engine { @@ -65,10 +65,6 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de bool is_padding = cfg.is_padding; bool is_dropout = (is_training && dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - auto bias_b = b; - auto bias_h = h; - auto bias_sq = s_q; - auto bias_skv = s_kv; NVTE_CHECK(~is_bias, "FP8 fused attention does not support pre/post_scale_bias yet!"); NVTE_CHECK(~is_alibi, "FP8 fused attention does not support ALiBi yet!"); bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && @@ -134,12 +130,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de // [SHARED-CACHE] Process-wide graph cache (was `static thread_local`) so a compiled graph // is reused across threads instead of rebuilt per thread. Safe because cuDNN >= 9.0 allows // concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe - // execute(); the static_asserts below fail the build loudly on an older toolkit. - static_assert(CUDNN_VERSION >= 91100, - "[SHARED-CACHE] shared fused-attn graph cache requires cuDNN >= 9.11 " - "(TE minimum supported cuDNN version)"); - static_assert(CUDNN_FRONTEND_VERSION >= 12500, - "[SHARED-CACHE] shared fused-attn graph cache requires cudnn-frontend >= 1.25.0"); + // execute(). The TE minimum-cuDNN-version bump that formalizes this requirement is a follow-up PR. static CacheType sdpa_fp8_fprop_cache; static std::mutex sdpa_fp8_fprop_cache_mutex; @@ -156,8 +147,8 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - fused_attn_graph_debug::note_cache_lookup("fwd", cache_hit, cfg); // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + graph_cache_debug::note_cache_lookup("fwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] + if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] return cached_graph; } @@ -408,29 +399,21 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] - GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); - GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); - GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) - GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->validate()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - fused_attn_graph_debug::note_fwd_build(); // [GRAPH-DEBUG] - if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] - std::vector serialized_graph; // [GRAPH-DEBUG] - if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] - fused_attn_graph_debug::note_graph_size("fwd", serialized_graph.size()); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] + graph_cache_debug::note_fwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - fused_attn_graph_debug::note_cache_size("fwd", cache.size()); // [GRAPH-DEBUG] - return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; } }; @@ -448,7 +431,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } - fused_attn_graph_debug::note_fwd_exec(); // [GRAPH-DEBUG] + graph_cache_debug::note_fwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -571,10 +554,6 @@ void fused_attn_fp8_bwd_impl( bool is_padding = cfg.is_padding; bool is_dropout = (dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - auto bias_b = b; - auto bias_h = h; - auto bias_sq = s_q; - auto bias_skv = s_kv; NVTE_CHECK(~is_bias, "FP8 fused attention does not support pre/post_scale_bias yet!"); NVTE_CHECK(~is_alibi, "FP8 fused attention does not support ALiBi yet!"); bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && @@ -659,8 +638,8 @@ void fused_attn_fp8_bwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - fused_attn_graph_debug::note_cache_lookup("bwd", cache_hit, cfg); // [GRAPH-DEBUG] - if (cache_hit && !fused_attn_graph_debug::cache_disabled()) { // [GRAPH-DEBUG] + graph_cache_debug::note_cache_lookup("bwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] + if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] return cached_graph; } @@ -1040,30 +1019,22 @@ void fused_attn_fp8_bwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - GRAPH_DEBUG_TIME_STAGE(Validate, NVTE_CHECK_CUDNN_FE(mha_graph->validate())); // [GRAPH-DEBUG] - GRAPH_DEBUG_TIME_STAGE(BuildOpGraph, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle))); - GRAPH_DEBUG_TIME_STAGE(CreatePlans, // [GRAPH-DEBUG] - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A}))); - GRAPH_DEBUG_TIME_STAGE(CheckSupport, NVTE_CHECK_CUDNN_FE(mha_graph->check_support())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) - GRAPH_DEBUG_TIME_STAGE(BuildPlans, NVTE_CHECK_CUDNN_FE(mha_graph->build_plans())); // [GRAPH-DEBUG] no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->validate()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - fused_attn_graph_debug::note_bwd_build(); // [GRAPH-DEBUG] - if (fused_attn_graph_debug::enabled()) { // [GRAPH-DEBUG] - std::vector serialized_graph; // [GRAPH-DEBUG] - if (mha_graph->serialize(serialized_graph).is_good()) // [GRAPH-DEBUG] - fused_attn_graph_debug::note_graph_size("bwd", serialized_graph.size()); // [GRAPH-DEBUG] - } // [GRAPH-DEBUG] + graph_cache_debug::note_bwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - fused_attn_graph_debug::note_cache_size("bwd", cache.size()); // [GRAPH-DEBUG] - return fused_attn_graph_debug::cache_disabled() ? return_tuple : inserted.first->second; + return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; } }; auto [mha_graph, Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, @@ -1080,7 +1051,7 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } - fused_attn_graph_debug::note_bwd_exec(); // [GRAPH-DEBUG] + graph_cache_debug::note_bwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1172,6 +1143,8 @@ void fused_attn_fp8_bwd_impl( } // namespace fused_attn +using namespace transformer_engine::fused_attn; + // fused attention FWD FP8 with separate Q, K, V void fused_attn_fp8_fwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, const Tensor* input_SoftmaxOffset, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index f75906cbad..4a408772e3 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,6 +8,9 @@ * \brief Functions for fused attention for FP8 */ +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ + #include #include "config_and_params.h" @@ -16,7 +19,7 @@ namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd(const FusedAttnConfig &cfg, const Tensor *input_Q, const Tensor *input_K, +void fused_attn_fp8_fwd(const fused_attn::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, @@ -24,7 +27,7 @@ void fused_attn_fp8_fwd(const FusedAttnConfig &cfg, const Tensor *input_Q, const cudnnHandle_t handle); // fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd(const FusedAttnConfig &cfg, const Tensor *input_Q, const Tensor *input_K, +void fused_attn_fp8_bwd(const fused_attn::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, @@ -36,10 +39,12 @@ void fused_attn_fp8_bwd(const FusedAttnConfig &cfg, const Tensor *input_Q, const // 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 explaining why it is not supported. -std::string is_supported_fp8_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); +std::string is_supported_fp8_fwd(const fused_attn::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 explaining why it is not supported. -std::string is_supported_fp8_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle); +std::string is_supported_fp8_bwd(const fused_attn::FusedAttnConfig &cfg, cudnnHandle_t handle); } // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h new file mode 100644 index 0000000000..d4464f1908 --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -0,0 +1,274 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// ============================================================================ +// Fused-attention graph-cache diagnostics. +// +// Lightweight, opt-in instrumentation for the cuDNN fused-attention graph cache. All output is +// gated behind an env switch and costs ~one cached-bool branch per fwd/bwd launch when off, so it +// is safe to leave compiled in for production. Enable at runtime with: +// export NVTE_FUSED_ATTN_CACHE_DEBUG=1 +// +// What it reports (all lines prefixed "[FUSED-ATTN-CACHE]"): +// - "BUILD" line whenever a new graph is constructed, plus a "SUMMARY" line at process exit with +// total graph builds vs. executions (fwd/bwd). Rebuilds >> executions => redundant construction +// (a make_cache_key / operator< that is missing a field, or a cache that is not being shared). +// - "HIT"/"MISS" line per cache lookup with the full (pre-normalization) config key, to diagnose +// stale-cache reuse: a HIT means two configs compared equal under operator<, so the field that +// distinguishes a wrongly-reused graph is one make_cache_key() normalized away or operator< +// omits -- diff a wrong HIT against the earlier BUILD to find it. +// - "thd ... path=legacy|direct" per THD (ragged) lookup and a "THD-PATH" summary, showing which +// impl path (bucketed batch vs. real cu_seqlens) the graph was built for. Low builds/lookups on +// the legacy path means batch bucketing is collapsing distinct batch sizes onto shared graphs. +// - Every line is tagged with a short per-thread id (tid=N) so cross-thread rebuilds of an +// identical key are visible. +// +// Separately, force every lookup to miss (never reuse a cached graph) with: +// export NVTE_FUSED_ATTN_DISABLE_CACHE=1 +// If a suite that fails with the cache enabled passes with it disabled, the bug is stale-cache +// reuse (an incomplete make_cache_key / operator<). +// +// ---------------------------------------------------------------------------- +// Reference numbers from earlier, heavier instrumentation (FE build-stage timing + cached-graph +// host-memory footprint), which was removed to keep this header lean. Collected by running +// tests/pytorch/attention/test_attention.py on GB200; each stage is invoked on the order of 2000 +// times over the run. +// +// FE build pipeline is dominated by build_plans() (cuDNN plan compilation / autotune): +// stage avg/call share of build cost +// validate 0.020 ms ~0% (was a static bool check previously) +// build_operation_graph 1.828 ms ~0.3% +// create_execution_plans 2.163 ms ~0.3% +// check_support 0.021 ms ~0% +// build_plans 618.673 ms >99% (dominates total build time) +// Note: avg/call is a full-suite mean; build_plans in particular scales with problem size and +// varies widely from call to call, so treat ~600 ms as an order-of-magnitude figure, not a +// constant. +// => The "real check_support" availability probe is essentially free; the entire expense is +// plan compilation, which only happens on a cache MISS. This is exactly what the graph cache + +// make_cache_key() normalization exist to avoid, so cache correctness (not probe cost) is what +// matters for performance. +// +// Cached-graph host memory (serialized graph size; a proxy for the plan/engine/tensor metadata +// each built graph holds -- device workspace is separate, sized per execute()): +// pass entries graphs avg/graph total +// fwd 670 1224 189.5 KB ~232 MB +// bwd 473 757 300.0 KB ~227 MB +// => ~190 KB (fwd) / ~300 KB (bwd) per distinct config; a long-lived process that sees many +// distinct shapes can accumulate hundreds of MB of cached graph metadata. Worth remembering if +// cache growth (rather than build time) ever becomes the concern. +// ---------------------------------------------------------------------------- +// ============================================================================ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ + +#include +#include +#include +#include + +#include "config_and_params.h" // for FusedAttnConfig field dump + +namespace transformer_engine { +namespace fused_attn { +namespace graph_cache_debug { + +// Short, stable per-thread id (0, 1, 2, ...) assigned on first use. Tagging every lookup with its +// thread id makes cross-thread rebuilds of an identical key visible. +inline unsigned thread_seq_id() { + static std::atomic next{0}; + static thread_local unsigned id = next.fetch_add(1); + return id; +} + +inline std::atomic &fwd_built() { + static std::atomic v{0}; + return v; +} +inline std::atomic &fwd_exec() { + static std::atomic v{0}; + return v; +} +inline std::atomic &bwd_built() { + static std::atomic v{0}; + return v; +} +inline std::atomic &bwd_exec() { + static std::atomic v{0}; + return v; +} + +// THD (ragged) cache lookups split by which impl path the graph was built for: +// legacy = batch quantized into a bucket (many batch sizes share one graph) +// direct = cu_seqlens fed to cuDNN directly (real batch baked in, no batch sharing) +// "builds" counts the lookups that actually constructed a new graph. A low builds/lookups ratio on +// the legacy path is the visible sign that batch bucketing is collapsing distinct batch sizes onto +// shared graphs. +inline std::atomic &thd_legacy_lookup() { + static std::atomic v{0}; + return v; +} +inline std::atomic &thd_legacy_build() { + static std::atomic v{0}; + return v; +} +inline std::atomic &thd_direct_lookup() { + static std::atomic v{0}; + return v; +} +inline std::atomic &thd_direct_build() { + static std::atomic v{0}; + return v; +} + +inline bool enabled() { + static const bool on = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return on; +} + +inline void dump(const char *event) { + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", + event, thread_seq_id(), static_cast(fwd_built().load()), + static_cast(fwd_exec().load()), + static_cast(bwd_built().load()), + static_cast(bwd_exec().load())); + std::fflush(stderr); +} + +inline void dump_thd_summary() { + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] THD-PATH | legacy lookups=%llu builds=%llu | direct lookups=%llu builds=%llu\n", + static_cast(thd_legacy_lookup().load()), + static_cast(thd_legacy_build().load()), + static_cast(thd_direct_lookup().load()), + static_cast(thd_direct_build().load())); + std::fflush(stderr); +} + +inline void register_summary_once() { + static const bool registered = [] { + std::atexit([] { + if (enabled()) { + dump("SUMMARY"); + dump_thd_summary(); + } + }); + return true; + }(); + (void)registered; +} + +inline void note_fwd_build() { + if (!enabled()) return; + register_summary_once(); + fwd_built().fetch_add(1); + dump("fwd BUILD"); +} +inline void note_fwd_exec() { + if (!enabled()) return; + register_summary_once(); + fwd_exec().fetch_add(1); +} +inline void note_bwd_build() { + if (!enabled()) return; + register_summary_once(); + bwd_built().fetch_add(1); + dump("bwd BUILD"); +} +inline void note_bwd_exec() { + if (!enabled()) return; + register_summary_once(); + bwd_exec().fetch_add(1); +} + +// Returns true when the graph cache should be bypassed (every lookup treated as a miss so a fresh +// graph is built each call). Gated by NVTE_FUSED_ATTN_DISABLE_CACHE. +inline bool cache_disabled() { + static const bool off = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_DISABLE_CACHE"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return off; +} + +// Logs one graph-cache lookup with its outcome (HIT/MISS) and the *real* (pre-normalization) config +// fields. A std::map HIT means the two configs compare equal under operator<, so the field that +// actually distinguishes a wrongly-reused graph is one that make_cache_key() normalized away or +// that operator< omits -- pass the real cfg (not the normalized cache key) here so that difference +// is visible when diffing a wrong HIT against the earlier BUILD that created the reused graph. +inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { + if (!enabled()) return; + register_summary_once(); + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld bias=%lld " + "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " + "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " + "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " + "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " + "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", + pass, hit ? "HIT" : "MISS", + (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", thread_seq_id(), + static_cast(c.is_training), + static_cast(c.deterministic), static_cast(c.cuda_graph), + static_cast(c.return_max_logit), static_cast(c.is_forward), + static_cast(c.attn_mask_type), static_cast(c.bias_type), + static_cast(c.window_size_left), static_cast(c.window_size_right), + static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), + static_cast(c.scaling_mode), static_cast(c.dropout), + static_cast(c.attn_scale), static_cast(c.qkv_dtype), + static_cast(c.o_dtype), static_cast(c.do_dtype), + static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), + static_cast(c.o_format), static_cast(c.do_format), + static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), + static_cast(c.do_scale_inv_format), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.head_dim_qk), static_cast(c.head_dim_v), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), + static_cast(c.num_pages_v), static_cast(c.page_size_k), + static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), + static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), + static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), + static_cast(c.bias_seqlen_kv)); + std::fflush(stderr); +} + +// Records, for one THD (ragged) cache lookup, which impl path the graph was built for -- +// "legacy" (batch quantized into a bucket) vs "direct" (real batch fed via cu_seqlens) -- and +// whether it hit the cache. `built` should reflect whether a new graph was actually constructed +// (i.e. a real miss, or a hit forced to rebuild by NVTE_FUSED_ATTN_DISABLE_CACHE). Comparing +// per-path lookups vs builds in the THD-PATH summary shows the batch-bucketing effect. +inline void note_thd_lookup(const char *pass, bool hit, bool built, bool legacy) { + if (!enabled()) return; + register_summary_once(); + if (legacy) { + thd_legacy_lookup().fetch_add(1); + if (built) thd_legacy_build().fetch_add(1); + } else { + thd_direct_lookup().fetch_add(1); + if (built) thd_direct_build().fetch_add(1); + } + std::fprintf(stderr, "[FUSED-ATTN-CACHE] thd %-3s %-4s | tid=%u | path=%s%s\n", pass, + hit ? "HIT" : "MISS", thread_seq_id(), legacy ? "legacy" : "direct", + (hit && built) ? " [cache-disabled->rebuild]" : ""); + std::fflush(stderr); +} + +} // namespace graph_cache_debug +} // namespace fused_attn +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ diff --git a/transformer_engine/common/fused_attn/graph_debug.h b/transformer_engine/common/fused_attn/graph_debug.h deleted file mode 100644 index 258a4556c4..0000000000 --- a/transformer_engine/common/fused_attn/graph_debug.h +++ /dev/null @@ -1,476 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -// ============================================================================ -// [GRAPH-DEBUG] TEMPORARY DEBUG INSTRUMENTATION -- REMOVE AFTER VERIFICATION. -// -// Counts fused-attention cuDNN graph *builds* (cache misses that construct a new -// graph) vs. *executions* (real forward/backward runs, excluding workspace-sizing -// probes) to detect redundant graph construction. Also logs every graph-cache -// lookup (HIT/MISS + the key fields) to diagnose stale-cache reuse across tests. -// -// Enable at runtime with: export NVTE_FUSED_ATTN_GRAPH_DEBUG=1 -// - A "BUILD" line is printed whenever a new graph is constructed. -// - A "HIT"/"MISS" line with the key fields is printed on every cache lookup. -// - A "thd ... path=legacy|direct" line is printed on every THD (ragged) lookup, showing -// which impl path (bucketed batch vs. real batch) the graph was built for. -// - A "SUMMARY" line with final build/exec totals is printed at process exit, followed by a -// "THD-PATH" line with per-path lookup/build totals (low builds/lookups on the legacy path -// means batch bucketing is collapsing distinct batch sizes onto shared graphs), and one -// "STAGE " line per FE build stage (validate ... build_plans) with total CPU/wall -// time and call count -- on `main` these were static boolean checks (~0 cost). -// -// Separately, force every lookup to miss (never reuse a cached graph) with: -// export NVTE_FUSED_ATTN_DISABLE_CACHE=1 -// If a suite that fails with the cache enabled passes with it disabled, the bug -// is stale-cache reuse (an incomplete make_cache_key / operator<). -// -// To remove all of this instrumentation later: -// 1. Delete this file (graph_debug.h). -// 2. Remove every line tagged with the "[GRAPH-DEBUG]" marker in: -// - fused_attn_fp8.cu -// - fused_attn_f16_arbitrary_seqlen.cu -// ============================================================================ - -#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ -#define TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// [GRAPH-DEBUG] Backtrace printing needs glibc's and libstdc++'s -// (demangling). Gate on availability so non-glibc toolchains still build; dump_backtrace() -// becomes a no-op there. -#if defined(__has_include) -#if __has_include() && __has_include() -#define NVTE_FUSED_ATTN_GRAPH_DEBUG_HAVE_BACKTRACE 1 -#include -#include -#endif -#endif - -#include "config_and_params.h" // [GRAPH-DEBUG] for FusedAttnConfig field dump - -namespace transformer_engine { -namespace fused_attn_graph_debug { - -// Short, stable per-thread id (0, 1, 2, ...) assigned on first use. The graph caches are -// static thread_local, so a graph built on one thread is invisible to another; tagging every -// lookup with its thread id makes cross-thread rebuilds of an identical key visible. -inline unsigned thread_seq_id() { - static std::atomic next{0}; - static thread_local unsigned id = next.fetch_add(1); - return id; -} - -inline std::atomic &fwd_built() { - static std::atomic v{0}; - return v; -} -inline std::atomic &fwd_exec() { - static std::atomic v{0}; - return v; -} -inline std::atomic &bwd_built() { - static std::atomic v{0}; - return v; -} -inline std::atomic &bwd_exec() { - static std::atomic v{0}; - return v; -} - -// THD (ragged) cache lookups split by which impl path the graph was built for: -// legacy = batch quantized into a bucket (many batch sizes share one graph) -// direct = cu_seqlens fed to cuDNN directly (real batch baked in, no batch sharing) -// "builds" counts the lookups that actually constructed a new graph. A low builds/lookups -// ratio on the legacy path is the visible sign that batch bucketing is collapsing distinct -// batch sizes onto shared graphs. -inline std::atomic &thd_legacy_lookup() { - static std::atomic v{0}; - return v; -} -inline std::atomic &thd_legacy_build() { - static std::atomic v{0}; - return v; -} -inline std::atomic &thd_direct_lookup() { - static std::atomic v{0}; - return v; -} -inline std::atomic &thd_direct_build() { - static std::atomic v{0}; - return v; -} - -inline bool enabled() { - static const bool on = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; - }(); - return on; -} - -inline void dump(const char *event) { - std::fprintf( - stderr, - "[GRAPH-DEBUG] %-10s | tid=%u | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", event, - thread_seq_id(), static_cast(fwd_built().load()), - static_cast(fwd_exec().load()), - static_cast(bwd_built().load()), - static_cast(bwd_exec().load())); - std::fflush(stderr); -} - -inline void dump_thd_summary() { - std::fprintf(stderr, - "[GRAPH-DEBUG] THD-PATH | legacy lookups=%llu builds=%llu | direct lookups=%llu " - "builds=%llu\n", - static_cast(thd_legacy_lookup().load()), - static_cast(thd_legacy_build().load()), - static_cast(thd_direct_lookup().load()), - static_cast(thd_direct_build().load())); - std::fflush(stderr); -} - -// [GRAPH-DEBUG] Host-memory footprint of cached graphs, split fwd/bwd (index 0/1). -// serialized bytes: size of fe::graph::Graph::serialize() output -- a proxy for the host -// memory one built graph holds (its plan / engine config / tensor metadata). Summed over -// builds; avg = sum / count gives the per-graph host cost. -// cache entries: high-water number of live graphs in the shared std::map (one graph per key). -// Device memory (workspace) is separate and sized per execute(), not held by the cached graph. -inline int pass_index(const char *pass) { return (pass[0] == 'b') ? 1 : 0; } // "bwd" -> 1 - -inline std::atomic &serial_bytes(int i) { - static std::array, 2> v{}; - return v[i]; -} -inline std::atomic &serial_count(int i) { - static std::array, 2> v{}; - return v[i]; -} -inline std::atomic &cache_entries(int i) { - static std::array, 2> v{}; - return v[i]; -} - -inline void dump_memory_summary() { - for (int i = 0; i < 2; ++i) { - const char *pass = (i == 0) ? "fwd" : "bwd"; - uint64_t cnt = serial_count(i).load(); - uint64_t bytes = serial_bytes(i).load(); - uint64_t entries = cache_entries(i).load(); - double total_kb = static_cast(bytes) / 1024.0; - double avg_kb = cnt ? total_kb / static_cast(cnt) : 0.0; - std::fprintf( - stderr, - "[GRAPH-DEBUG] MEMORY %-3s | cache entries=%llu | serialized graphs=%llu total=%.1f KB (avg %.1f KB)\n", - pass, static_cast(entries), static_cast(cnt), - total_kb, avg_kb); - } - std::fflush(stderr); -} - -// [GRAPH-DEBUG] Per-stage CPU/wall time for the FE build pipeline (validate ... build_plans). -// On `main` these were static boolean checks (~0 cost); this quantifies the added cost. -enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; - -inline const char *stage_name(BuildStage s) { - switch (s) { - case BuildStage::Validate: - return "validate"; - case BuildStage::BuildOpGraph: - return "build_operation_graph"; - case BuildStage::CreatePlans: - return "create_execution_plans"; - case BuildStage::CheckSupport: - return "check_support"; - case BuildStage::BuildPlans: - return "build_plans"; - default: - return "?"; - } -} - -inline std::atomic &stage_calls(BuildStage s) { - static std::array, static_cast(BuildStage::kCount)> v{}; - return v[static_cast(s)]; -} -inline std::atomic &stage_cpu_ns(BuildStage s) { - static std::array, static_cast(BuildStage::kCount)> v{}; - return v[static_cast(s)]; -} -inline std::atomic &stage_wall_ns(BuildStage s) { - static std::array, static_cast(BuildStage::kCount)> v{}; - return v[static_cast(s)]; -} - -inline void dump_stage_summary() { - for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { - BuildStage s = static_cast(i); - uint64_t n = stage_calls(s).load(); - if (n == 0) continue; - double cpu_ms = static_cast(stage_cpu_ns(s).load()) / 1e6; - double wall_ms = static_cast(stage_wall_ns(s).load()) / 1e6; - std::fprintf(stderr, - "[GRAPH-DEBUG] STAGE %-22s | calls=%llu | cpu=%.1f ms (avg %.3f ms) | wall=%.1f ms\n", - stage_name(s), static_cast(n), cpu_ms, cpu_ms / n, wall_ms); - } - std::fflush(stderr); -} - -inline void register_summary_once() { - static const bool registered = [] { - std::atexit([] { - if (enabled()) { - dump("SUMMARY"); - dump_thd_summary(); - dump_stage_summary(); - dump_memory_summary(); - } - }); - return true; - }(); - (void)registered; -} - -inline void note_fwd_build() { - if (!enabled()) return; - register_summary_once(); - fwd_built().fetch_add(1); - dump("fwd BUILD"); -} -inline void note_fwd_exec() { - if (!enabled()) return; - register_summary_once(); - fwd_exec().fetch_add(1); -} -inline void note_bwd_build() { - if (!enabled()) return; - register_summary_once(); - bwd_built().fetch_add(1); - dump("bwd BUILD"); -} -inline void note_bwd_exec() { - if (!enabled()) return; - register_summary_once(); - bwd_exec().fetch_add(1); -} - -// [GRAPH-DEBUG] Record the serialized size (host-memory proxy) of one freshly built graph. -// Call only after a successful serialize() so the average reflects real graphs. -inline void note_graph_size(const char *pass, size_t serialized_bytes) { - if (!enabled()) return; - register_summary_once(); - int i = pass_index(pass); - serial_bytes(i).fetch_add(serialized_bytes); - serial_count(i).fetch_add(1); -} - -// [GRAPH-DEBUG] Record the current shared-cache entry count (kept as a high-water mark). -inline void note_cache_size(const char *pass, size_t entries) { - if (!enabled()) return; - register_summary_once(); - int i = pass_index(pass); - uint64_t prev = cache_entries(i).load(); - while (entries > prev && !cache_entries(i).compare_exchange_weak(prev, entries)) { - } -} - -// Returns true when the graph cache should be bypassed (every lookup treated as a -// miss so a fresh graph is built each call). Gated by NVTE_FUSED_ATTN_DISABLE_CACHE. -inline bool cache_disabled() { - static const bool off = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_DISABLE_CACHE"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; - }(); - return off; -} - -// [GRAPH-DEBUG] Opt-in C++ backtrace printing next to each cache lookup. Kept separate from the -// main NVTE_FUSED_ATTN_GRAPH_DEBUG switch because a full stack per lookup is very verbose; enable -// with NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE=1 (the main switch must also be on). -inline bool backtrace_enabled() { - static const bool on = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; - }(); - return on; -} - -// [GRAPH-DEBUG] Frames to print per lookup (override with NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE_DEPTH). -inline int backtrace_depth() { - static const int depth = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG_BACKTRACE_DEPTH"); - int d = (e != nullptr && e[0] != '\0') ? std::atoi(e) : 24; - if (d < 1) d = 1; - if (d > 128) d = 128; - return d; - }(); - return depth; -} - -// [GRAPH-DEBUG] Print a symbolized (and, where possible, demangled) C++ backtrace, one frame per -// line, each tagged so the frames group visually under the HIT/MISS line they belong to. `skip` -// drops the top frames that are just this instrumentation (dump_backtrace + its caller). For -// readable function names the library must be built/linked with -rdynamic (or -g); otherwise -// non-exported frames show as "(+0x)". -inline void dump_backtrace(const char *tag, int skip = 2) { - if (!backtrace_enabled()) return; -#if defined(NVTE_FUSED_ATTN_GRAPH_DEBUG_HAVE_BACKTRACE) - const int max_frames = backtrace_depth() + skip; - std::vector frames(static_cast(max_frames)); - int n = ::backtrace(frames.data(), max_frames); - if (n <= skip) return; - char **symbols = ::backtrace_symbols(frames.data(), n); - if (symbols == nullptr) return; - for (int i = skip; i < n; ++i) { - // glibc format: "(+0x) [0x]"; demangle the "" span. - std::string line = symbols[i]; - char *open = std::strchr(symbols[i], '('); - char *plus = open ? std::strchr(open, '+') : nullptr; - if (open != nullptr && plus != nullptr && plus > open + 1) { - std::string mangled(open + 1, plus); - int status = 0; - char *demangled = abi::__cxa_demangle(mangled.c_str(), nullptr, nullptr, &status); - if (status == 0 && demangled != nullptr) { - line = std::string(symbols[i], open + 1) + demangled + plus; - std::free(demangled); - } - } - std::fprintf(stderr, "[GRAPH-DEBUG] bt[%-4s] #%02d %s\n", tag, i - skip, line.c_str()); - } - std::fflush(stderr); - std::free(symbols); -#else - (void)tag; - (void)skip; -#endif -} - -// Logs one graph-cache lookup with its outcome (HIT/MISS) and the *real* (pre- -// normalization) config fields. A std::map HIT means the two configs compare equal -// under operator<, so the field that actually distinguishes a wrongly-reused graph -// is one that make_cache_key() normalized away or that operator< omits -- pass the -// real cfg (not the normalized cache key) here so that difference is visible when -// diffing a wrong HIT against the earlier BUILD that created the reused graph. -inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { - if (!enabled()) return; - register_summary_once(); - std::fprintf( - stderr, - "[GRAPH-DEBUG] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld " - "bias=%lld " - "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " - "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " - "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " - "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " - "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", - pass, hit ? "HIT" : "MISS", (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", - thread_seq_id(), static_cast(c.is_training), static_cast(c.deterministic), - static_cast(c.cuda_graph), static_cast(c.return_max_logit), - static_cast(c.is_forward), static_cast(c.attn_mask_type), - static_cast(c.bias_type), static_cast(c.window_size_left), - static_cast(c.window_size_right), static_cast(c.bottom_right_diagonal), - static_cast(c.softmax_type), static_cast(c.scaling_mode), - static_cast(c.dropout), static_cast(c.attn_scale), - static_cast(c.qkv_dtype), static_cast(c.o_dtype), - static_cast(c.do_dtype), static_cast(c.dqkv_dtype), - static_cast(c.qkv_layout), static_cast(c.o_format), - static_cast(c.do_format), static_cast(c.dqkv_layout), - static_cast(c.qkv_scale_inv_format), static_cast(c.do_scale_inv_format), - static_cast(c.batch_size), static_cast(c.num_attn_heads), - static_cast(c.num_gqa_groups), static_cast(c.head_dim_qk), - static_cast(c.head_dim_v), static_cast(c.max_seqlen_q), - static_cast(c.max_seqlen_kv), static_cast(c.num_tokens_q), - static_cast(c.num_tokens_kv), static_cast(c.bucketed_batch_size), - static_cast(c.bucketed_num_tokens_q), - static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), - static_cast(c.num_pages_v), static_cast(c.page_size_k), - static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), - static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), - static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), - static_cast(c.bias_seqlen_kv)); - std::fflush(stderr); - dump_backtrace(hit ? "HIT" : "MISS"); // [GRAPH-DEBUG] frames for this fwd/bwd lookup -} - -// Records, for one THD (ragged) cache lookup, which impl path the graph was built for -- -// "legacy" (batch quantized into a bucket) vs "direct" (real batch fed via cu_seqlens) -- and -// whether it hit the cache. `built` should reflect whether a new graph was actually constructed -// (i.e. a real miss, or a hit forced to rebuild by NVTE_FUSED_ATTN_DISABLE_CACHE). Comparing -// per-path lookups vs builds in the THD-PATH summary shows the batch-bucketing effect. -inline void note_thd_lookup(const char *pass, bool hit, bool built, bool legacy) { - if (!enabled()) return; - register_summary_once(); - if (legacy) { - thd_legacy_lookup().fetch_add(1); - if (built) thd_legacy_build().fetch_add(1); - } else { - thd_direct_lookup().fetch_add(1); - if (built) thd_direct_build().fetch_add(1); - } - std::fprintf(stderr, "[GRAPH-DEBUG] thd %-3s %-4s | tid=%u | path=%s%s\n", pass, - hit ? "HIT" : "MISS", thread_seq_id(), legacy ? "legacy" : "direct", - (hit && built) ? " [cache-disabled->rebuild]" : ""); - std::fflush(stderr); - dump_backtrace(hit ? "HIT" : "MISS"); // [GRAPH-DEBUG] frames for this THD (ragged) lookup -} - -// [GRAPH-DEBUG] Thread-CPU clock (excludes time blocked on locks / GPU sync), in nanoseconds. -inline uint64_t cpu_now_ns() { - timespec ts; - clock_gettime(CLOCK_THREAD_CPUTIME_ID, &ts); - return static_cast(ts.tv_sec) * 1000000000ull + static_cast(ts.tv_nsec); -} - -// [GRAPH-DEBUG] RAII timer: records wall + thread-CPU time for one FE build stage. Zero cost -// (only an enabled() bool read) when NVTE_FUSED_ATTN_GRAPH_DEBUG is unset. The destructor records -// even on early return / thrown NVTE_CHECK, so it is safe to wrap the checked FE calls. -struct ScopedStageTimer { - BuildStage stage; - bool on; - std::chrono::steady_clock::time_point w0; - uint64_t c0{0}; - explicit ScopedStageTimer(BuildStage s) : stage(s), on(enabled()) { - if (!on) return; - register_summary_once(); - c0 = cpu_now_ns(); - w0 = std::chrono::steady_clock::now(); - } - ~ScopedStageTimer() { - if (!on) return; - uint64_t cpu = cpu_now_ns() - c0; - uint64_t wall = static_cast( - std::chrono::duration_cast(std::chrono::steady_clock::now() - w0) - .count()); - stage_cpu_ns(stage).fetch_add(cpu); - stage_wall_ns(stage).fetch_add(wall); - stage_calls(stage).fetch_add(1); - } -}; - -} // namespace fused_attn_graph_debug -} // namespace transformer_engine - -// [GRAPH-DEBUG] Wrap a single (possibly NVTE_CHECK_*-guarded) FE call to time it under `stage`. -#define GRAPH_DEBUG_TIME_STAGE(stage, expr) \ - do { \ - ::transformer_engine::fused_attn_graph_debug::ScopedStageTimer _gd_stage_timer( \ - ::transformer_engine::fused_attn_graph_debug::BuildStage::stage); \ - expr; \ - } while (0) - -#endif // TRANSFORMER_ENGINE_FUSED_ATTN_GRAPH_DEBUG_H_ diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 875ccdbe72..c6c8957b1b 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -8,7 +8,6 @@ #include #include "../common.h" -#include "../cudnn_utils.h" #include "../util/cuda_runtime.h" #include "transformer_engine/fused_attn.h" #include "utils.h" @@ -325,93 +324,6 @@ void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int6 } } -bool allowAllConfig(cudnnBackendDescriptor_t engine_config) { - (void)engine_config; - return false; -} - -cudnn_frontend::Tensor tensor_create(cudnnDataType_t type, int64_t id, int64_t const *dim, - int64_t const *stride, bool is_virtual, bool is_value) { - int nbDims = 4; - auto tensor_created = - cudnn_frontend::TensorBuilder() - .setDim(nbDims, dim) - .setStride(nbDims, stride) - .setId(id) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(type) - .setVirtual(is_virtual) - .setByValue(is_value) - .build(); - return tensor_created; -} - -cudnn_frontend::Tensor tensor_create_with_offset( - cudnnDataType_t type, int64_t id, int64_t const *dim, int64_t const *stride, bool is_virtual, - bool is_value, std::shared_ptr raggedOffset) { - int nbDims = 4; - auto tensor_created = - cudnn_frontend::TensorBuilder() - .setDim(nbDims, dim) - .setStride(nbDims, stride) - .setId(id) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(type) - .setVirtual(is_virtual) - .setByValue(is_value) - .setRaggedOffset(raggedOffset) - .build(); - return tensor_created; -} - -cudnn_frontend::PointWiseDesc pw_desc_create(cudnnDataType_t type, cudnnPointwiseMode_t mode) { - auto pw_desc_created = - cudnn_frontend::PointWiseDescBuilder().setMode(mode).setComputeType(type).build(); - return pw_desc_created; -} - -cudnn_frontend::Operation unary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc) { - auto pw_op_created = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR) - .setxDesc(xDesc) - .setyDesc(yDesc) - .setpwDesc(pwDesc) - .build(); - return pw_op_created; -} - -cudnn_frontend::Operation binary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc) { - auto pw_op_created = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR) - .setxDesc(xDesc) - .setbDesc(bDesc) - .setyDesc(yDesc) - .setpwDesc(pwDesc) - .build(); - return pw_op_created; -} - -cudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &tDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc) { - auto pw_op_created = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR) - .setxDesc(xDesc) - .setbDesc(bDesc) - .settDesc(tDesc) - .setyDesc(yDesc) - .setpwDesc(pwDesc) - .build(); - return pw_op_created; -} - // convert cu_seqlens to actual_seqlens __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 7ac78ee4e1..e240e2a421 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -4,12 +4,8 @@ * See LICENSE for license information. ************************************************************************/ -#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_UTILS_H_ -#define TRANSFORMER_ENGINE_FUSED_ATTN_UTILS_H_ - -#include -#include -#include +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ #include #include @@ -223,32 +219,6 @@ inline void generateMatrixStridesWithLayout(int64_t b, int64_t h, int64_t hg, in void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix); -bool allowAllConfig(cudnnBackendDescriptor_t engine_config); - -cudnn_frontend::Tensor tensor_create(cudnnDataType_t type, int64_t id, int64_t const *dim, - int64_t const *stride, bool is_virtual, bool is_value); - -cudnn_frontend::Tensor tensor_create_with_offset( - cudnnDataType_t type, int64_t id, int64_t const *dim, int64_t const *stride, bool is_virtual, - bool is_value, std::shared_ptr raggedOffset); - -cudnn_frontend::PointWiseDesc pw_desc_create(cudnnDataType_t type, cudnnPointwiseMode_t mode); - -cudnn_frontend::Operation unary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc); - -cudnn_frontend::Operation binary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc); - -cudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &tDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc); - // Per-tensor scale factors relating cu_seqlens_padded (token units) to tensor-element // ragged offsets, as a function of the QKV layout group. Single source of truth shared // by the cu_seqlens_padded_to_offsets conversion kernel and the direct-seqlens path @@ -334,4 +304,4 @@ uint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t len, cud } // namespace fused_attn } // namespace transformer_engine -#endif +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ 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 8c2e181eb2..609e4ef550 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 @@ -63,7 +63,14 @@ # Setup Attention Logging attn_log.setup_logging() -# Global vars for available attention backends and ALiBi cache +# Global vars for available attention backends and ALiBi cache. +# +# `_attention_backends` holds the most-recently-selected backend result plus the +# `backend_selection_requires_update` flag. The flag is the public invalidation signal: external +# callers (e.g. the test suite) set it to True to force a full re-selection, typically because they +# changed an NVTE_* environment toggle that is not captured by AttentionParams. This dict is kept +# for backward compatibility (its shape and the flag are part of the de-facto public API); the +# actual multi-entry caching lives in `_attention_backend_cache` below. _attention_backends = { "attention_params": None, "use_flash_attention": None, @@ -74,6 +81,87 @@ "backend_selection_requires_update": False, } +# LRU cache of backend-selection results, so that alternating between a handful of configs in the +# same run does not repay get_attention_backend() on every switch (the previous single-slot cache +# thrashed whenever two or more configs interleaved). AttentionParams is unhashable -- it holds +# dicts/lists/tensors and its custom __eq__ disables __hash__ -- so we cannot use it as a dict key. +# Instead we keep an insertion-ordered list of {"attention_params", "env_key", } and +# linear-scan. Capacity is small (10), so the scan is negligible next to a real selection. +# +# The cache identity is (env_key, attention_params). env_key captures the NVTE_* environment toggles +# that get_attention_backend() reads at call time but that AttentionParams does not encode. Including +# it means flipping any such toggle naturally misses and re-selects, so callers do NOT need to +# manually invalidate after changing the environment. Setting +# _attention_backends["backend_selection_requires_update"] = True still hard-clears the whole cache +# for anyone who wants to start completely afresh (e.g. after changing GPU/arch mid-process). +_ATTENTION_BACKEND_RESULT_KEYS = ( + "use_flash_attention", + "flash_attention_backend", + "use_fused_attention", + "fused_attention_backend", + "use_unfused_attention", +) +_ATTENTION_BACKEND_CACHE_MAXSIZE = 10 +_attention_backend_cache = [] + +# Explicit allow-list of the NVTE_* toggles that steer get_attention_backend(). We deliberately do +# NOT snapshot the whole NVTE_* namespace: unrelated toggles (determinism, non-attention modules like +# Linear, debug/logging, etc.) would otherwise needlessly invalidate cached selections. +# +# IMPORTANT: keep this in sync with the os.getenv(...) reads inside +# dot_product_attention/utils.py::get_attention_backend(). If that function begins consulting a new +# NVTE_* toggle that is not listed here, the cache can return a stale (wrong) backend selection. +_ATTENTION_BACKEND_ENV_VARS = ( + "NVTE_FLASH_ATTN", + "NVTE_FLASH_ATTN_V2", + "NVTE_FLASH_ATTN_V3", + "NVTE_FLASH_ATTN_V4", + "NVTE_FUSED_ATTN", + "NVTE_UNFUSED_ATTN", + "NVTE_FP8_DPA_BWD", + "NVTE_DPA_FP8CS_O_in_F16", + "NVTE_DPA_FP8_RECIPE", + "NVTE_DPA_FP8_FORMAT", + "NVTE_DPA_FP8DS_AMAX_ALGO", + "NVTE_DPA_FP8DS_AMAX_HISTLEN", + "NVTE_DPA_FP8DS_REDUCE_AMAX", + "NVTE_UnfusedDPA_Emulate_FP8", +) + + +def _attention_env_key(): + """Snapshot of the selection-relevant NVTE_* toggles (see _ATTENTION_BACKEND_ENV_VARS). + + These influence get_attention_backend() but are not captured by AttentionParams, so they must be + part of the cache identity to avoid returning a result computed under a different environment. A + value of None means the variable is unset (i.e. get_attention_backend() would use its default). + """ + return tuple(os.environ.get(name) for name in _ATTENTION_BACKEND_ENV_VARS) + + +def _attention_backend_cache_lookup(attention_params, env_key): + """Return the cached result dict matching ``(env_key, attention_params)`` (promoted to MRU).""" + for i, entry in enumerate(_attention_backend_cache): + # Compare the cheap env_key tuple before the per-field AttentionParams.__eq__. + if entry["env_key"] == env_key and entry["attention_params"] == attention_params: + if i != len(_attention_backend_cache) - 1: + _attention_backend_cache.append(_attention_backend_cache.pop(i)) + return entry + return None + + +def _attention_backend_cache_store(attention_params, env_key, result): + """Insert/refresh the entry for ``(env_key, attention_params)`` as MRU and evict beyond capacity.""" + for i, entry in enumerate(_attention_backend_cache): + if entry["env_key"] == env_key and entry["attention_params"] == attention_params: + _attention_backend_cache.pop(i) + break + entry = {"attention_params": attention_params, "env_key": env_key, **result} + _attention_backend_cache.append(entry) + while len(_attention_backend_cache) > _ATTENTION_BACKEND_CACHE_MAXSIZE: + _attention_backend_cache.pop(0) + return entry + _alibi_cache = { "_num_heads": None, "_alibi_slopes": None, @@ -1039,8 +1127,8 @@ def forward( .. note:: Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`, - and :attr:`NVTE_FUSED_ATTN_BACKEND` to control which DotProductAttention backend, - and FusedAttention backend if applicable, to use. Transformer Engine first filters + and :attr:`NVTE_UNFUSED_ATTN` to control which DotProductAttention backend to use. + Transformer Engine first filters backends by support for the runtime environment and input configuration, then applies a performance-based preference order. On supported pre-Hopper GPUs, FlashAttention is preferred over FusedAttention and UnfusedDotProductAttention when both optimized @@ -1635,13 +1723,17 @@ def forward( use_fused_attention = False use_unfused_attention = True else: - if ( - _attention_backends["attention_params"] is None - or attention_params != _attention_backends["attention_params"] - ): - _attention_backends["attention_params"] = attention_params - _attention_backends["backend_selection_requires_update"] = True + # A forced update hard-clears the entire cache. This is optional now that the + # cache identity includes the NVTE_* environment (so env changes miss on their own); + # it remains as an explicit "start completely afresh" hook (e.g. after changing + # GPU/arch mid-process) for callers who want it. if _attention_backends["backend_selection_requires_update"]: + _attention_backend_cache.clear() + _attention_backends["backend_selection_requires_update"] = False + + env_key = _attention_env_key() + cached = _attention_backend_cache_lookup(attention_params, env_key) + if cached is None: ( use_flash_attention, flash_attention_backend, @@ -1650,14 +1742,17 @@ def forward( use_unfused_attention, _, ) = dpa_utils.get_attention_backend(attention_params) - # Set global _attention_backends var using return value - # from get_attention_backend() - _attention_backends["use_flash_attention"] = use_flash_attention - _attention_backends["flash_attention_backend"] = flash_attention_backend - _attention_backends["use_fused_attention"] = use_fused_attention - _attention_backends["fused_attention_backend"] = fused_attention_backend - _attention_backends["use_unfused_attention"] = use_unfused_attention - _attention_backends["backend_selection_requires_update"] = False + cached = _attention_backend_cache_store( + attention_params, + env_key, + { + "use_flash_attention": use_flash_attention, + "flash_attention_backend": flash_attention_backend, + "use_fused_attention": use_fused_attention, + "fused_attention_backend": fused_attention_backend, + "use_unfused_attention": use_unfused_attention, + }, + ) if use_flash_attention: self.logger.info( "Running with FlashAttention backend (version %s)", @@ -1671,11 +1766,17 @@ def forward( elif use_unfused_attention: self.logger.info("Running with UnfusedDotProductAttention backend") else: - use_flash_attention = _attention_backends["use_flash_attention"] - flash_attention_backend = _attention_backends["flash_attention_backend"] - use_fused_attention = _attention_backends["use_fused_attention"] - fused_attention_backend = _attention_backends["fused_attention_backend"] - use_unfused_attention = _attention_backends["use_unfused_attention"] + use_flash_attention = cached["use_flash_attention"] + flash_attention_backend = cached["flash_attention_backend"] + use_fused_attention = cached["use_fused_attention"] + fused_attention_backend = cached["fused_attention_backend"] + use_unfused_attention = cached["use_unfused_attention"] + + # Mirror the active selection into the legacy single-slot dict so its public shape + # (and any external readers) keep working as before. + _attention_backends["attention_params"] = attention_params + for _key in _ATTENTION_BACKEND_RESULT_KEYS: + _attention_backends[_key] = cached[_key] # raise exception if no backend is available if sum([use_flash_attention, use_fused_attention, use_unfused_attention]) == 0: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py b/transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py deleted file mode 100644 index b69839f74c..0000000000 --- a/transformer_engine/pytorch/attention/dot_product_attention/graph_debug.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -# ============================================================================ -# [GRAPH-DEBUG] TEMPORARY DEBUG INSTRUMENTATION -- REMOVE AFTER VERIFICATION. -# -# Python-side companion to the C++ instrumentation in -# common/fused_attn/graph_debug.h. Prints the Python call stack that leads into -# each fused-attention backend query / forward / backward call, so the Python -# frames interleave (on stderr) just above the C++ "[GRAPH-DEBUG] fwd/bwd HIT|MISS" -# lines they trigger. This makes it possible to attribute each cuDNN graph-cache -# lookup to the exact Python caller (availability probe vs. module backend -# re-selection vs. actual fwd/bwd execution). -# -# Enable with the SAME switch as the C++ side: -# export NVTE_FUSED_ATTN_GRAPH_DEBUG=1 -# Optionally cap the number of printed frames (default 12): -# export NVTE_FUSED_ATTN_GRAPH_DEBUG_PY_DEPTH= -# -# To remove all of this instrumentation later: -# 1. Delete this file (graph_debug.py). -# 2. Remove every line tagged with the "[GRAPH-DEBUG]" marker in: -# - attention/dot_product_attention/utils.py -# - cpp_extensions/fused_attn.py -# ============================================================================ - -import os -import sys -import threading -import traceback - -_enabled = None -_depth = None - - -def enabled(): - """True when NVTE_FUSED_ATTN_GRAPH_DEBUG is set (same switch as the C++ side).""" - global _enabled - if _enabled is None: - val = os.getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG", "") - _enabled = val not in ("", "0") - return _enabled - - -def _depth_val(): - global _depth - if _depth is None: - val = os.getenv("NVTE_FUSED_ATTN_GRAPH_DEBUG_PY_DEPTH", "") - try: - _depth = int(val) if val else 12 - except ValueError: - _depth = 12 - _depth = max(1, min(_depth, 128)) - return _depth - - -def pytrace(tag): - """Print a compact Python call stack to stderr, tagged so it groups with the C++ - [GRAPH-DEBUG] frames that follow. No-op unless NVTE_FUSED_ATTN_GRAPH_DEBUG is set.""" - if not enabled(): - return - # Drop this frame (pytrace itself); show the most recent frames, oldest first. - frames = traceback.extract_stack()[:-1][-_depth_val() :] - out = sys.stderr - out.write(f"[GRAPH-DEBUG-PY] {tag} | tid={threading.get_ident()}\n") - for fr in frames: - code = f" -> {fr.line}" if fr.line else "" - out.write(f"[GRAPH-DEBUG-PY] {fr.filename}:{fr.lineno} {fr.name}(){code}\n") - out.flush() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 76685874aa..694038752d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -426,11 +426,6 @@ def get_attention_backend( All available backends that could support the provided input. A list of Booleans in the form of [use_flash_attention, use_fused_attention, use_unfused_attention]. """ - # [GRAPH-DEBUG] Trace the Python caller that triggers a fused-attn backend query (maps to the - # C++ support-check "fwd/bwd HIT|MISS" lines from is_supported_f16_*). Remove after verification. - from transformer_engine.pytorch.attention.dot_product_attention import graph_debug - - graph_debug.pytrace("get_attention_backend") # NOTE: As part of refactoring attention.py, populating the _attention_backends cache in attention # is no longer performed at the end of get_attention_backend(), but the responsibility of doing so # is shifted over to the caller of this function diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 3e68eab85b..9c22c56bd1 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -303,12 +303,6 @@ def fused_attn_fwd( else: raise ValueError(f"Unsupported backend {fused_attention_backend}") - # [GRAPH-DEBUG] Trace the Python caller of the actual fwd kernel (maps to the C++ execution - # "fwd HIT|MISS" line + note_fwd_exec). Remove after verification. - from transformer_engine.pytorch.attention.dot_product_attention import graph_debug - - graph_debug.pytrace("fused_attn_fwd (execute)") - # execute kernel output_tensors = tex.fused_attn_fwd( max_seqlen_q, @@ -559,12 +553,6 @@ def fused_attn_bwd( f" for backend={fused_attention_backend}." ) - # [GRAPH-DEBUG] Trace the Python caller of the actual bwd kernel (maps to the C++ execution - # "bwd HIT|MISS" line + note_bwd_exec). Remove after verification. - from transformer_engine.pytorch.attention.dot_product_attention import graph_debug - - graph_debug.pytrace("fused_attn_bwd (execute)") - output_tensors = tex.fused_attn_bwd( max_seqlen_q, max_seqlen_kv, From aa34ccb356772e4854e8fc930a8d104d39cbcd0a Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 06:44:29 -0700 Subject: [PATCH 37/63] use macros for attr_sizes[], cache_key_tuple(), and fprintf in cache debug Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.h | 151 ++++++++++-------- .../common/fused_attn/graph_cache_debug.h | 46 ++---- 2 files changed, 98 insertions(+), 99 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index ab01de9b91..34560ddc93 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -19,6 +19,64 @@ namespace transformer_engine { namespace fused_attn { +// Single source of truth for the graph-cache-relevant fields of FusedAttnConfig, in the SAME order +// as NVTEFusedAttnConfigAttribute / attr_sizes[]. Each row is +// X(member, wire_type, printf_fmt, printf_cast, debug_label) +// where wire_type is the field's 1-byte-exact serialization type (bool is serialized as uint8_t), +// and printf_fmt / printf_cast / debug_label drive the [FUSED-ATTN-CACHE] debug dump. operator<, +// attr_sizes[], and that debug dump are all generated from this one list so they cannot drift apart. +// When adding/removing a cache-relevant field, edit ONLY this list -- and the public +// NVTEFusedAttnConfigAttribute enum, whose entry count the static_assert below cross-checks. +#define TE_FUSED_ATTN_CACHE_KEY_FIELDS(X) \ + /* basic attention settings */ \ + X(is_training, uint8_t, "%d", int, "train") \ + X(deterministic, uint8_t, "%d", int, "det") \ + X(cuda_graph, uint8_t, "%d", int, "cg") \ + X(return_max_logit, uint8_t, "%d", int, "maxlogit") \ + X(attn_mask_type, NVTE_Mask_Type, "%lld", long long, "mask") \ + X(bias_type, NVTE_Bias_Type, "%lld", long long, "bias") \ + X(window_size_left, int64_t, "%lld", long long, "wl") \ + X(window_size_right, int64_t, "%lld", long long, "wr") \ + X(bottom_right_diagonal, uint8_t, "%d", int, "brd") \ + X(softmax_type, NVTE_Softmax_Type, "%lld", long long, "softmax") \ + X(scaling_mode, NVTEScalingMode, "%lld", long long, "scale_mode") \ + X(dropout, float, "%g", double, "dropout") \ + X(attn_scale, float, "%g", double, "attn_scale") \ + /* tensor types */ \ + X(qkv_dtype, NVTEDType, "%lld", long long, "qkv_dt") \ + X(o_dtype, NVTEDType, "%lld", long long, "o_dt") \ + X(do_dtype, NVTEDType, "%lld", long long, "do_dt") \ + X(dqkv_dtype, NVTEDType, "%lld", long long, "dqkv_dt") \ + /* tensor layouts */ \ + X(qkv_layout, NVTE_QKV_Layout, "%lld", long long, "qkv_lay") \ + X(o_format, NVTE_QKV_Format, "%lld", long long, "o_fmt") \ + X(do_format, NVTE_QKV_Format, "%lld", long long, "do_fmt") \ + X(dqkv_layout, NVTE_QKV_Layout, "%lld", long long, "dqkv_lay") \ + X(qkv_scale_inv_format, NVTE_QKV_Format, "%lld", long long, "qkv_sif") \ + X(do_scale_inv_format, NVTE_QKV_Format, "%lld", long long, "do_sif") \ + /* tensor dimensions */ \ + X(batch_size, size_t, "%lld", long long, "b") \ + X(num_attn_heads, size_t, "%lld", long long, "h") \ + X(num_gqa_groups, size_t, "%lld", long long, "hg") \ + X(head_dim_qk, size_t, "%lld", long long, "dqk") \ + X(head_dim_v, size_t, "%lld", long long, "dv") \ + X(max_seqlen_q, size_t, "%lld", long long, "sq") \ + X(max_seqlen_kv, size_t, "%lld", long long, "skv") \ + X(num_tokens_q, size_t, "%lld", long long, "tq") \ + X(num_tokens_kv, size_t, "%lld", long long, "tkv") \ + /* paged KV dimensions */ \ + X(num_pages_k, size_t, "%lld", long long, "npk") \ + X(num_pages_v, size_t, "%lld", long long, "npv") \ + X(page_size_k, size_t, "%lld", long long, "psk") \ + X(page_size_v, size_t, "%lld", long long, "psv") \ + X(max_pages_per_seq_k, size_t, "%lld", long long, "mppk") \ + X(max_pages_per_seq_v, size_t, "%lld", long long, "mppv") \ + /* bias dimensions */ \ + X(bias_batch_size, size_t, "%lld", long long, "bias_b") \ + X(bias_num_heads, size_t, "%lld", long long, "bias_h") \ + X(bias_seqlen_q, size_t, "%lld", long long, "bias_sq") \ + X(bias_seqlen_kv, size_t, "%lld", long long, "bias_skv") + struct FusedAttnConfig { // basic attention settings bool is_training = true; @@ -96,78 +154,27 @@ struct FusedAttnConfig { bool is_causal = false; bool is_causal_bottom_right = false; + // Generated from TE_FUSED_ATTN_CACHE_KEY_FIELDS so the per-attribute serialized sizes stay in lockstep + // with the field list (and, via the static_assert below, with NVTEFusedAttnConfigAttribute). static constexpr size_t attr_sizes[] = { - // basic attention settings - sizeof(uint8_t), // is_training - sizeof(uint8_t), // deterministic - sizeof(uint8_t), // cuda_graph - sizeof(uint8_t), // return_max_logit - sizeof(NVTE_Mask_Type), // attn_mask_type - sizeof(NVTE_Bias_Type), // bias_type - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(NVTEScalingMode), // scaling_mode - sizeof(float), // dropout - sizeof(float), // attn_scale - // tensor types - sizeof(NVTEDType), // qkv_dtype - sizeof(NVTEDType), // o_dtype - sizeof(NVTEDType), // do_dtype - sizeof(NVTEDType), // dqkv_dtype - // tensor layouts - 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 - // tensor dimensions - 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), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - sizeof(size_t), // num_tokens_q - sizeof(size_t), // num_tokens_kv - // paged KV dimensions - 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 - // bias dimensions - sizeof(size_t), // bias_batch_size - sizeof(size_t), // bias_num_heads - sizeof(size_t), // bias_seqlen_q - sizeof(size_t), // bias_seqlen_kv +#define TE_FUSED_ATTN_CACHE_KEY_FIELD_SIZE(member, wire, fmt, cast, label) sizeof(wire), + TE_FUSED_ATTN_CACHE_KEY_FIELDS(TE_FUSED_ATTN_CACHE_KEY_FIELD_SIZE) +#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_SIZE }; + // Tuple of all cache-relevant fields, generated from TE_FUSED_ATTN_CACHE_KEY_FIELDS. The trailing 0 + // sentinel absorbs the macro's trailing comma; it is identical on both operands so it never + // affects ordering. Used by operator< so the comparison can never omit a field. + auto cache_key_tuple() const { + return std::make_tuple( +#define TE_FUSED_ATTN_CACHE_KEY_FIELD_VALUE(member, wire, fmt, cast, label) member, + TE_FUSED_ATTN_CACHE_KEY_FIELDS(TE_FUSED_ATTN_CACHE_KEY_FIELD_VALUE) +#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_VALUE + 0); + } + bool operator<(const FusedAttnConfig &rhs) const { - return std::tie(is_training, deterministic, cuda_graph, return_max_logit, attn_mask_type, - bias_type, window_size_left, window_size_right, bottom_right_diagonal, - softmax_type, scaling_mode, dropout, attn_scale, qkv_dtype, o_dtype, do_dtype, - dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, - head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_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_batch_size, bias_num_heads, bias_seqlen_q, - bias_seqlen_kv) < - std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, - rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, - rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, - rhs.attn_scale, rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, - rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, - rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, - rhs.num_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, - rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, - 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); + return cache_key_tuple() < rhs.cache_key_tuple(); } // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields @@ -181,6 +188,14 @@ struct FusedAttnConfig { FusedAttnConfig make_cache_key() const; }; +// Cross-check the generated field list against the public attribute enum: if a cache-relevant field +// is added to TE_FUSED_ATTN_CACHE_KEY_FIELDS without a matching NVTEFusedAttnConfigAttribute entry (or +// vice versa), this fails to compile instead of silently corrupting attribute (de)serialization. +static_assert(sizeof(FusedAttnConfig::attr_sizes) / sizeof(FusedAttnConfig::attr_sizes[0]) == + kNVTEFusedAttnConfigNumAttributes, + "TE_FUSED_ATTN_CACHE_KEY_FIELDS is out of sync with NVTEFusedAttnConfigAttribute; " + "update both together."); + inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); return reinterpret_cast(config); diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index d4464f1908..577c573891 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -209,40 +209,24 @@ inline bool cache_disabled() { inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { if (!enabled()) return; register_summary_once(); + // The cache-key portion of this line is generated from TE_FUSED_ATTN_CACHE_KEY_FIELDS, so it can never + // drift from operator< / attr_sizes. The internal-only fields that are NOT part of the cache key + // (is_forward, and the bucketed THD counts that make_cache_key() folds away) are appended + // explicitly at the end -- they are still worth printing to diagnose a wrongly-reused graph. +#define TE_FUSED_ATTN_CACHE_KEY_FIELD_FMT(member, wire, fmt, cast, label) label "=" fmt " " +#define TE_FUSED_ATTN_CACHE_KEY_FIELD_ARG(member, wire, fmt, cast, label) , static_cast(c.member) std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld bias=%lld " - "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " - "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " - "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " - "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " - "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", + "[FUSED-ATTN-CACHE] %-3s %-4s%s | tid=%u | " TE_FUSED_ATTN_CACHE_KEY_FIELDS( + TE_FUSED_ATTN_CACHE_KEY_FIELD_FMT) "fwd=%d bb=%lld btq=%lld btkv=%lld\n", pass, hit ? "HIT" : "MISS", - (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", thread_seq_id(), - static_cast(c.is_training), - static_cast(c.deterministic), static_cast(c.cuda_graph), - static_cast(c.return_max_logit), static_cast(c.is_forward), - static_cast(c.attn_mask_type), static_cast(c.bias_type), - static_cast(c.window_size_left), static_cast(c.window_size_right), - static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), - static_cast(c.scaling_mode), static_cast(c.dropout), - static_cast(c.attn_scale), static_cast(c.qkv_dtype), - static_cast(c.o_dtype), static_cast(c.do_dtype), - static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), - static_cast(c.o_format), static_cast(c.do_format), - static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), - static_cast(c.do_scale_inv_format), static_cast(c.batch_size), - static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), - static_cast(c.head_dim_qk), static_cast(c.head_dim_v), - static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), - static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), - static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), - static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), - static_cast(c.num_pages_v), static_cast(c.page_size_k), - static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), - static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), - static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), - static_cast(c.bias_seqlen_kv)); + (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", + thread_seq_id() TE_FUSED_ATTN_CACHE_KEY_FIELDS(TE_FUSED_ATTN_CACHE_KEY_FIELD_ARG), + static_cast(c.is_forward), static_cast(c.bucketed_batch_size), + static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv)); +#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_FMT +#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_ARG std::fflush(stderr); } From 642b58a0f09824a17cd6e8f9d13e0b088a2aabf8 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 06:44:45 -0700 Subject: [PATCH 38/63] Revert "use macros for attr_sizes[], cache_key_tuple(), and fprintf in cache debug" This reverts commit aa34ccb356772e4854e8fc930a8d104d39cbcd0a. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.h | 151 ++++++++---------- .../common/fused_attn/graph_cache_debug.h | 46 ++++-- 2 files changed, 99 insertions(+), 98 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 34560ddc93..ab01de9b91 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -19,64 +19,6 @@ namespace transformer_engine { namespace fused_attn { -// Single source of truth for the graph-cache-relevant fields of FusedAttnConfig, in the SAME order -// as NVTEFusedAttnConfigAttribute / attr_sizes[]. Each row is -// X(member, wire_type, printf_fmt, printf_cast, debug_label) -// where wire_type is the field's 1-byte-exact serialization type (bool is serialized as uint8_t), -// and printf_fmt / printf_cast / debug_label drive the [FUSED-ATTN-CACHE] debug dump. operator<, -// attr_sizes[], and that debug dump are all generated from this one list so they cannot drift apart. -// When adding/removing a cache-relevant field, edit ONLY this list -- and the public -// NVTEFusedAttnConfigAttribute enum, whose entry count the static_assert below cross-checks. -#define TE_FUSED_ATTN_CACHE_KEY_FIELDS(X) \ - /* basic attention settings */ \ - X(is_training, uint8_t, "%d", int, "train") \ - X(deterministic, uint8_t, "%d", int, "det") \ - X(cuda_graph, uint8_t, "%d", int, "cg") \ - X(return_max_logit, uint8_t, "%d", int, "maxlogit") \ - X(attn_mask_type, NVTE_Mask_Type, "%lld", long long, "mask") \ - X(bias_type, NVTE_Bias_Type, "%lld", long long, "bias") \ - X(window_size_left, int64_t, "%lld", long long, "wl") \ - X(window_size_right, int64_t, "%lld", long long, "wr") \ - X(bottom_right_diagonal, uint8_t, "%d", int, "brd") \ - X(softmax_type, NVTE_Softmax_Type, "%lld", long long, "softmax") \ - X(scaling_mode, NVTEScalingMode, "%lld", long long, "scale_mode") \ - X(dropout, float, "%g", double, "dropout") \ - X(attn_scale, float, "%g", double, "attn_scale") \ - /* tensor types */ \ - X(qkv_dtype, NVTEDType, "%lld", long long, "qkv_dt") \ - X(o_dtype, NVTEDType, "%lld", long long, "o_dt") \ - X(do_dtype, NVTEDType, "%lld", long long, "do_dt") \ - X(dqkv_dtype, NVTEDType, "%lld", long long, "dqkv_dt") \ - /* tensor layouts */ \ - X(qkv_layout, NVTE_QKV_Layout, "%lld", long long, "qkv_lay") \ - X(o_format, NVTE_QKV_Format, "%lld", long long, "o_fmt") \ - X(do_format, NVTE_QKV_Format, "%lld", long long, "do_fmt") \ - X(dqkv_layout, NVTE_QKV_Layout, "%lld", long long, "dqkv_lay") \ - X(qkv_scale_inv_format, NVTE_QKV_Format, "%lld", long long, "qkv_sif") \ - X(do_scale_inv_format, NVTE_QKV_Format, "%lld", long long, "do_sif") \ - /* tensor dimensions */ \ - X(batch_size, size_t, "%lld", long long, "b") \ - X(num_attn_heads, size_t, "%lld", long long, "h") \ - X(num_gqa_groups, size_t, "%lld", long long, "hg") \ - X(head_dim_qk, size_t, "%lld", long long, "dqk") \ - X(head_dim_v, size_t, "%lld", long long, "dv") \ - X(max_seqlen_q, size_t, "%lld", long long, "sq") \ - X(max_seqlen_kv, size_t, "%lld", long long, "skv") \ - X(num_tokens_q, size_t, "%lld", long long, "tq") \ - X(num_tokens_kv, size_t, "%lld", long long, "tkv") \ - /* paged KV dimensions */ \ - X(num_pages_k, size_t, "%lld", long long, "npk") \ - X(num_pages_v, size_t, "%lld", long long, "npv") \ - X(page_size_k, size_t, "%lld", long long, "psk") \ - X(page_size_v, size_t, "%lld", long long, "psv") \ - X(max_pages_per_seq_k, size_t, "%lld", long long, "mppk") \ - X(max_pages_per_seq_v, size_t, "%lld", long long, "mppv") \ - /* bias dimensions */ \ - X(bias_batch_size, size_t, "%lld", long long, "bias_b") \ - X(bias_num_heads, size_t, "%lld", long long, "bias_h") \ - X(bias_seqlen_q, size_t, "%lld", long long, "bias_sq") \ - X(bias_seqlen_kv, size_t, "%lld", long long, "bias_skv") - struct FusedAttnConfig { // basic attention settings bool is_training = true; @@ -154,27 +96,78 @@ struct FusedAttnConfig { bool is_causal = false; bool is_causal_bottom_right = false; - // Generated from TE_FUSED_ATTN_CACHE_KEY_FIELDS so the per-attribute serialized sizes stay in lockstep - // with the field list (and, via the static_assert below, with NVTEFusedAttnConfigAttribute). static constexpr size_t attr_sizes[] = { -#define TE_FUSED_ATTN_CACHE_KEY_FIELD_SIZE(member, wire, fmt, cast, label) sizeof(wire), - TE_FUSED_ATTN_CACHE_KEY_FIELDS(TE_FUSED_ATTN_CACHE_KEY_FIELD_SIZE) -#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_SIZE + // basic attention settings + sizeof(uint8_t), // is_training + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(NVTEScalingMode), // scaling_mode + sizeof(float), // dropout + sizeof(float), // attn_scale + // tensor types + sizeof(NVTEDType), // qkv_dtype + sizeof(NVTEDType), // o_dtype + sizeof(NVTEDType), // do_dtype + sizeof(NVTEDType), // dqkv_dtype + // tensor layouts + 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 + // tensor dimensions + 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), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(size_t), // num_tokens_q + sizeof(size_t), // num_tokens_kv + // paged KV dimensions + 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 + // bias dimensions + sizeof(size_t), // bias_batch_size + sizeof(size_t), // bias_num_heads + sizeof(size_t), // bias_seqlen_q + sizeof(size_t), // bias_seqlen_kv }; - // Tuple of all cache-relevant fields, generated from TE_FUSED_ATTN_CACHE_KEY_FIELDS. The trailing 0 - // sentinel absorbs the macro's trailing comma; it is identical on both operands so it never - // affects ordering. Used by operator< so the comparison can never omit a field. - auto cache_key_tuple() const { - return std::make_tuple( -#define TE_FUSED_ATTN_CACHE_KEY_FIELD_VALUE(member, wire, fmt, cast, label) member, - TE_FUSED_ATTN_CACHE_KEY_FIELDS(TE_FUSED_ATTN_CACHE_KEY_FIELD_VALUE) -#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_VALUE - 0); - } - bool operator<(const FusedAttnConfig &rhs) const { - return cache_key_tuple() < rhs.cache_key_tuple(); + return std::tie(is_training, deterministic, cuda_graph, return_max_logit, attn_mask_type, + bias_type, window_size_left, window_size_right, bottom_right_diagonal, + softmax_type, scaling_mode, dropout, attn_scale, qkv_dtype, o_dtype, do_dtype, + dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, + head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_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_batch_size, bias_num_heads, bias_seqlen_q, + bias_seqlen_kv) < + std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, + rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, + rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, + rhs.attn_scale, rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, + rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, + rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, + rhs.num_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, + rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, + 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); } // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields @@ -188,14 +181,6 @@ struct FusedAttnConfig { FusedAttnConfig make_cache_key() const; }; -// Cross-check the generated field list against the public attribute enum: if a cache-relevant field -// is added to TE_FUSED_ATTN_CACHE_KEY_FIELDS without a matching NVTEFusedAttnConfigAttribute entry (or -// vice versa), this fails to compile instead of silently corrupting attribute (de)serialization. -static_assert(sizeof(FusedAttnConfig::attr_sizes) / sizeof(FusedAttnConfig::attr_sizes[0]) == - kNVTEFusedAttnConfigNumAttributes, - "TE_FUSED_ATTN_CACHE_KEY_FIELDS is out of sync with NVTEFusedAttnConfigAttribute; " - "update both together."); - inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); return reinterpret_cast(config); diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 577c573891..d4464f1908 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -209,24 +209,40 @@ inline bool cache_disabled() { inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { if (!enabled()) return; register_summary_once(); - // The cache-key portion of this line is generated from TE_FUSED_ATTN_CACHE_KEY_FIELDS, so it can never - // drift from operator< / attr_sizes. The internal-only fields that are NOT part of the cache key - // (is_forward, and the bucketed THD counts that make_cache_key() folds away) are appended - // explicitly at the end -- they are still worth printing to diagnose a wrongly-reused graph. -#define TE_FUSED_ATTN_CACHE_KEY_FIELD_FMT(member, wire, fmt, cast, label) label "=" fmt " " -#define TE_FUSED_ATTN_CACHE_KEY_FIELD_ARG(member, wire, fmt, cast, label) , static_cast(c.member) std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %-3s %-4s%s | tid=%u | " TE_FUSED_ATTN_CACHE_KEY_FIELDS( - TE_FUSED_ATTN_CACHE_KEY_FIELD_FMT) "fwd=%d bb=%lld btq=%lld btkv=%lld\n", + "[FUSED-ATTN-CACHE] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld bias=%lld " + "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " + "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " + "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " + "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " + "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", pass, hit ? "HIT" : "MISS", - (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", - thread_seq_id() TE_FUSED_ATTN_CACHE_KEY_FIELDS(TE_FUSED_ATTN_CACHE_KEY_FIELD_ARG), - static_cast(c.is_forward), static_cast(c.bucketed_batch_size), - static_cast(c.bucketed_num_tokens_q), - static_cast(c.bucketed_num_tokens_kv)); -#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_FMT -#undef TE_FUSED_ATTN_CACHE_KEY_FIELD_ARG + (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", thread_seq_id(), + static_cast(c.is_training), + static_cast(c.deterministic), static_cast(c.cuda_graph), + static_cast(c.return_max_logit), static_cast(c.is_forward), + static_cast(c.attn_mask_type), static_cast(c.bias_type), + static_cast(c.window_size_left), static_cast(c.window_size_right), + static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), + static_cast(c.scaling_mode), static_cast(c.dropout), + static_cast(c.attn_scale), static_cast(c.qkv_dtype), + static_cast(c.o_dtype), static_cast(c.do_dtype), + static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), + static_cast(c.o_format), static_cast(c.do_format), + static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), + static_cast(c.do_scale_inv_format), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.head_dim_qk), static_cast(c.head_dim_v), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), + static_cast(c.num_pages_v), static_cast(c.page_size_k), + static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), + static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), + static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), + static_cast(c.bias_seqlen_kv)); std::fflush(stderr); } From 97cb2b7467d97ae60027e66e3039cf5c5fcce5de Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 07:05:24 -0700 Subject: [PATCH 39/63] remove graph cache debug code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 12 - .../fused_attn_f16_arbitrary_seqlen.cu | 30 +- .../common/fused_attn/fused_attn_fp8.cu | 19 +- .../common/fused_attn/graph_cache_debug.h | 274 ------------------ 4 files changed, 11 insertions(+), 324 deletions(-) delete mode 100644 transformer_engine/common/fused_attn/graph_cache_debug.h diff --git a/docs/envvars.rst b/docs/envvars.rst index e543975f59..bf32df8971 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -177,18 +177,6 @@ backend-selection overview. :Default: ``0`` :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. -.. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG - - :Type: ``int`` (0 or 1) - :Default: ``0`` - :Description: Enable diagnostic logging for the cuDNN FusedAttention graph cache. When set to ``1``, prints to stderr (prefixed ``[FUSED-ATTN-CACHE]``) a per-lookup HIT/MISS line with the full graph-cache key, a BUILD line whenever a new graph is constructed, and a SUMMARY of graph builds vs. executions at process exit. Useful for diagnosing redundant graph rebuilds or stale-cache reuse. Has negligible overhead when unset. - -.. envvar:: NVTE_FUSED_ATTN_DISABLE_CACHE - - :Type: ``int`` (0 or 1) - :Default: ``0`` - :Description: Bypass the cuDNN FusedAttention graph cache, rebuilding a fresh graph on every forward/backward call. Intended for debugging stale-cache reuse only: it forces expensive graph recompilation on every call and must not be used in production. If a run that fails with the cache enabled passes with it disabled, the bug is stale-cache reuse (an incomplete cache key). Pairs with :envvar:`NVTE_FUSED_ATTN_CACHE_DEBUG` for inspecting each rebuild. - .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO :Type: ``int`` (0 or 1) 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 b0316b83ff..176773fb51 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 @@ -10,7 +10,7 @@ #include #include -#include // [SHARED-CACHE] +#include #include #include "../common.h" @@ -18,7 +18,6 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" -#include "graph_cache_debug.h" // [FUSED-ATTN-CACHE] #include "utils.h" namespace transformer_engine { @@ -174,14 +173,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - graph_cache_debug::note_cache_lookup("fwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [FUSED-ATTN-CACHE] - sm_arch_ != 120) { // [FUSED-ATTN-CACHE] - graph_cache_debug::note_thd_lookup( // [FUSED-ATTN-CACHE] - "fwd", cache_hit, !cache_hit || graph_cache_debug::cache_disabled(), - /*legacy=*/!use_cu_seqlens_directly); // [FUSED-ATTN-CACHE] - } // [FUSED-ATTN-CACHE] - if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] + if (cache_hit) { return cached_graph; } @@ -458,13 +450,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - graph_cache_debug::note_fwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; + return inserted.first->second; } }; @@ -499,7 +490,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - graph_cache_debug::note_fwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -730,15 +720,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - graph_cache_debug::note_cache_lookup("bwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600 && // [FUSED-ATTN-CACHE] - sm_arch_ != 120) { // [FUSED-ATTN-CACHE] - // The backward impl has no cu_seqlens-directly path; it always buckets the batch. - graph_cache_debug::note_thd_lookup( // [FUSED-ATTN-CACHE] - "bwd", cache_hit, !cache_hit || graph_cache_debug::cache_disabled(), - /*legacy=*/true); // [FUSED-ATTN-CACHE] - } // [FUSED-ATTN-CACHE] - if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] + if (cache_hit) { return cached_graph; } @@ -986,13 +968,12 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - graph_cache_debug::note_bwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; + return inserted.first->second; } }; @@ -1022,7 +1003,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - graph_cache_debug::note_bwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index a98a2f1950..682af81f55 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -4,14 +4,13 @@ * See LICENSE for license information. ************************************************************************/ -#include // [SHARED-CACHE] -#include // [FUSED-ATTN-CACHE] serialized-size probe +#include +#include #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" -#include "graph_cache_debug.h" // [FUSED-ATTN-CACHE] #include "utils.h" namespace transformer_engine { @@ -147,8 +146,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - graph_cache_debug::note_cache_lookup("fwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] - if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] + if (cache_hit) { return cached_graph; } @@ -407,13 +405,12 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - graph_cache_debug::note_fwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; + return inserted.first->second; } }; @@ -431,7 +428,6 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } - graph_cache_debug::note_fwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -638,8 +634,7 @@ void fused_attn_fp8_bwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } - graph_cache_debug::note_cache_lookup("bwd", cache_hit, cfg); // [FUSED-ATTN-CACHE] - if (cache_hit && !graph_cache_debug::cache_disabled()) { // [FUSED-ATTN-CACHE] + if (cache_hit) { return cached_graph; } @@ -1028,13 +1023,12 @@ void fused_attn_fp8_bwd_impl( auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - graph_cache_debug::note_bwd_build(); // [FUSED-ATTN-CACHE] // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, // reuse theirs and discard ours so all threads share one graph (rare duplicate build). { std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); - return graph_cache_debug::cache_disabled() ? return_tuple : inserted.first->second; + return inserted.first->second; } }; auto [mha_graph, Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, @@ -1051,7 +1045,6 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } - graph_cache_debug::note_bwd_exec(); // [FUSED-ATTN-CACHE] // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h deleted file mode 100644 index d4464f1908..0000000000 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ /dev/null @@ -1,274 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -// ============================================================================ -// Fused-attention graph-cache diagnostics. -// -// Lightweight, opt-in instrumentation for the cuDNN fused-attention graph cache. All output is -// gated behind an env switch and costs ~one cached-bool branch per fwd/bwd launch when off, so it -// is safe to leave compiled in for production. Enable at runtime with: -// export NVTE_FUSED_ATTN_CACHE_DEBUG=1 -// -// What it reports (all lines prefixed "[FUSED-ATTN-CACHE]"): -// - "BUILD" line whenever a new graph is constructed, plus a "SUMMARY" line at process exit with -// total graph builds vs. executions (fwd/bwd). Rebuilds >> executions => redundant construction -// (a make_cache_key / operator< that is missing a field, or a cache that is not being shared). -// - "HIT"/"MISS" line per cache lookup with the full (pre-normalization) config key, to diagnose -// stale-cache reuse: a HIT means two configs compared equal under operator<, so the field that -// distinguishes a wrongly-reused graph is one make_cache_key() normalized away or operator< -// omits -- diff a wrong HIT against the earlier BUILD to find it. -// - "thd ... path=legacy|direct" per THD (ragged) lookup and a "THD-PATH" summary, showing which -// impl path (bucketed batch vs. real cu_seqlens) the graph was built for. Low builds/lookups on -// the legacy path means batch bucketing is collapsing distinct batch sizes onto shared graphs. -// - Every line is tagged with a short per-thread id (tid=N) so cross-thread rebuilds of an -// identical key are visible. -// -// Separately, force every lookup to miss (never reuse a cached graph) with: -// export NVTE_FUSED_ATTN_DISABLE_CACHE=1 -// If a suite that fails with the cache enabled passes with it disabled, the bug is stale-cache -// reuse (an incomplete make_cache_key / operator<). -// -// ---------------------------------------------------------------------------- -// Reference numbers from earlier, heavier instrumentation (FE build-stage timing + cached-graph -// host-memory footprint), which was removed to keep this header lean. Collected by running -// tests/pytorch/attention/test_attention.py on GB200; each stage is invoked on the order of 2000 -// times over the run. -// -// FE build pipeline is dominated by build_plans() (cuDNN plan compilation / autotune): -// stage avg/call share of build cost -// validate 0.020 ms ~0% (was a static bool check previously) -// build_operation_graph 1.828 ms ~0.3% -// create_execution_plans 2.163 ms ~0.3% -// check_support 0.021 ms ~0% -// build_plans 618.673 ms >99% (dominates total build time) -// Note: avg/call is a full-suite mean; build_plans in particular scales with problem size and -// varies widely from call to call, so treat ~600 ms as an order-of-magnitude figure, not a -// constant. -// => The "real check_support" availability probe is essentially free; the entire expense is -// plan compilation, which only happens on a cache MISS. This is exactly what the graph cache + -// make_cache_key() normalization exist to avoid, so cache correctness (not probe cost) is what -// matters for performance. -// -// Cached-graph host memory (serialized graph size; a proxy for the plan/engine/tensor metadata -// each built graph holds -- device workspace is separate, sized per execute()): -// pass entries graphs avg/graph total -// fwd 670 1224 189.5 KB ~232 MB -// bwd 473 757 300.0 KB ~227 MB -// => ~190 KB (fwd) / ~300 KB (bwd) per distinct config; a long-lived process that sees many -// distinct shapes can accumulate hundreds of MB of cached graph metadata. Worth remembering if -// cache growth (rather than build time) ever becomes the concern. -// ---------------------------------------------------------------------------- -// ============================================================================ - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ - -#include -#include -#include -#include - -#include "config_and_params.h" // for FusedAttnConfig field dump - -namespace transformer_engine { -namespace fused_attn { -namespace graph_cache_debug { - -// Short, stable per-thread id (0, 1, 2, ...) assigned on first use. Tagging every lookup with its -// thread id makes cross-thread rebuilds of an identical key visible. -inline unsigned thread_seq_id() { - static std::atomic next{0}; - static thread_local unsigned id = next.fetch_add(1); - return id; -} - -inline std::atomic &fwd_built() { - static std::atomic v{0}; - return v; -} -inline std::atomic &fwd_exec() { - static std::atomic v{0}; - return v; -} -inline std::atomic &bwd_built() { - static std::atomic v{0}; - return v; -} -inline std::atomic &bwd_exec() { - static std::atomic v{0}; - return v; -} - -// THD (ragged) cache lookups split by which impl path the graph was built for: -// legacy = batch quantized into a bucket (many batch sizes share one graph) -// direct = cu_seqlens fed to cuDNN directly (real batch baked in, no batch sharing) -// "builds" counts the lookups that actually constructed a new graph. A low builds/lookups ratio on -// the legacy path is the visible sign that batch bucketing is collapsing distinct batch sizes onto -// shared graphs. -inline std::atomic &thd_legacy_lookup() { - static std::atomic v{0}; - return v; -} -inline std::atomic &thd_legacy_build() { - static std::atomic v{0}; - return v; -} -inline std::atomic &thd_direct_lookup() { - static std::atomic v{0}; - return v; -} -inline std::atomic &thd_direct_build() { - static std::atomic v{0}; - return v; -} - -inline bool enabled() { - static const bool on = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; - }(); - return on; -} - -inline void dump(const char *event) { - std::fprintf( - stderr, - "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%llu exec=%llu | bwd built=%llu exec=%llu\n", - event, thread_seq_id(), static_cast(fwd_built().load()), - static_cast(fwd_exec().load()), - static_cast(bwd_built().load()), - static_cast(bwd_exec().load())); - std::fflush(stderr); -} - -inline void dump_thd_summary() { - std::fprintf( - stderr, - "[FUSED-ATTN-CACHE] THD-PATH | legacy lookups=%llu builds=%llu | direct lookups=%llu builds=%llu\n", - static_cast(thd_legacy_lookup().load()), - static_cast(thd_legacy_build().load()), - static_cast(thd_direct_lookup().load()), - static_cast(thd_direct_build().load())); - std::fflush(stderr); -} - -inline void register_summary_once() { - static const bool registered = [] { - std::atexit([] { - if (enabled()) { - dump("SUMMARY"); - dump_thd_summary(); - } - }); - return true; - }(); - (void)registered; -} - -inline void note_fwd_build() { - if (!enabled()) return; - register_summary_once(); - fwd_built().fetch_add(1); - dump("fwd BUILD"); -} -inline void note_fwd_exec() { - if (!enabled()) return; - register_summary_once(); - fwd_exec().fetch_add(1); -} -inline void note_bwd_build() { - if (!enabled()) return; - register_summary_once(); - bwd_built().fetch_add(1); - dump("bwd BUILD"); -} -inline void note_bwd_exec() { - if (!enabled()) return; - register_summary_once(); - bwd_exec().fetch_add(1); -} - -// Returns true when the graph cache should be bypassed (every lookup treated as a miss so a fresh -// graph is built each call). Gated by NVTE_FUSED_ATTN_DISABLE_CACHE. -inline bool cache_disabled() { - static const bool off = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_DISABLE_CACHE"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; - }(); - return off; -} - -// Logs one graph-cache lookup with its outcome (HIT/MISS) and the *real* (pre-normalization) config -// fields. A std::map HIT means the two configs compare equal under operator<, so the field that -// actually distinguishes a wrongly-reused graph is one that make_cache_key() normalized away or -// that operator< omits -- pass the real cfg (not the normalized cache key) here so that difference -// is visible when diffing a wrong HIT against the earlier BUILD that created the reused graph. -inline void note_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { - if (!enabled()) return; - register_summary_once(); - std::fprintf( - stderr, - "[FUSED-ATTN-CACHE] %-3s %-4s%s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld bias=%lld " - "wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " - "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " - "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " - "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " - "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", - pass, hit ? "HIT" : "MISS", - (hit && cache_disabled()) ? " [cache-disabled->rebuild]" : "", thread_seq_id(), - static_cast(c.is_training), - static_cast(c.deterministic), static_cast(c.cuda_graph), - static_cast(c.return_max_logit), static_cast(c.is_forward), - static_cast(c.attn_mask_type), static_cast(c.bias_type), - static_cast(c.window_size_left), static_cast(c.window_size_right), - static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), - static_cast(c.scaling_mode), static_cast(c.dropout), - static_cast(c.attn_scale), static_cast(c.qkv_dtype), - static_cast(c.o_dtype), static_cast(c.do_dtype), - static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), - static_cast(c.o_format), static_cast(c.do_format), - static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), - static_cast(c.do_scale_inv_format), static_cast(c.batch_size), - static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), - static_cast(c.head_dim_qk), static_cast(c.head_dim_v), - static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), - static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), - static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), - static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), - static_cast(c.num_pages_v), static_cast(c.page_size_k), - static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), - static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), - static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), - static_cast(c.bias_seqlen_kv)); - std::fflush(stderr); -} - -// Records, for one THD (ragged) cache lookup, which impl path the graph was built for -- -// "legacy" (batch quantized into a bucket) vs "direct" (real batch fed via cu_seqlens) -- and -// whether it hit the cache. `built` should reflect whether a new graph was actually constructed -// (i.e. a real miss, or a hit forced to rebuild by NVTE_FUSED_ATTN_DISABLE_CACHE). Comparing -// per-path lookups vs builds in the THD-PATH summary shows the batch-bucketing effect. -inline void note_thd_lookup(const char *pass, bool hit, bool built, bool legacy) { - if (!enabled()) return; - register_summary_once(); - if (legacy) { - thd_legacy_lookup().fetch_add(1); - if (built) thd_legacy_build().fetch_add(1); - } else { - thd_direct_lookup().fetch_add(1); - if (built) thd_direct_build().fetch_add(1); - } - std::fprintf(stderr, "[FUSED-ATTN-CACHE] thd %-3s %-4s | tid=%u | path=%s%s\n", pass, - hit ? "HIT" : "MISS", thread_seq_id(), legacy ? "legacy" : "direct", - (hit && built) ? " [cache-disabled->rebuild]" : ""); - std::fflush(stderr); -} - -} // namespace graph_cache_debug -} // namespace fused_attn -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ From 17e5fe0286f474d999baa1f4a0a595ab1e0d9183 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 08:17:59 -0700 Subject: [PATCH 40/63] review and clean up Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn_score_mod.py | 5 +- tests/pytorch/attention/test_attention.py | 17 ++ tests/pytorch/utils.py | 4 - .../common/fused_attn/config_and_params.cpp | 10 +- .../common/fused_attn/config_and_params.h | 4 +- .../common/fused_attn/fused_attn.cpp | 24 ++- .../fused_attn_f16_arbitrary_seqlen.cu | 36 ++-- .../common/fused_attn/fused_attn_fp8.cu | 34 ++-- .../include/transformer_engine/fused_attn.h | 36 ++-- transformer_engine/jax/attention.py | 4 - .../jax/cpp_extensions/attention.py | 6 +- transformer_engine/jax/flax/transformer.py | 5 +- .../dot_product_attention.py | 154 +++--------------- .../attention/dot_product_attention/utils.py | 8 +- 14 files changed, 117 insertions(+), 230 deletions(-) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 6de133f822..e965a08665 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -404,7 +404,10 @@ def __init__(self, *args, **kwargs): def get_fused_attn_backend(self): if kernel_available: return NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, "" - return NVTE_Fused_Attn_Backend.NVTE_No_Backend, "fake: no backend" + return ( + NVTE_Fused_Attn_Backend.NVTE_No_Backend, + "fake FusedAttnHelper: no fused attention backend available for this configuration", + ) def fake_fused_attn( qkv, diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 731958ec43..e63b7b7b04 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -23,6 +23,9 @@ is_fp8_available, is_bf16_available, ) +from transformer_engine.pytorch.attention.dot_product_attention import ( + _attention_backends, +) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, check_set_window_size, @@ -1025,6 +1028,8 @@ def _run_dot_product_attention( os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + # Create seqlens qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) if "padding" in config.attn_mask_type or qkv_format == "thd": @@ -1579,6 +1584,8 @@ def _run_transformer_layer( os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + # Create input tensor if qkv_format == "sbhd": inp = torch.randn( @@ -2033,6 +2040,7 @@ def test_mha_fp8_vs_f16( os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") flash_attn_fwd_fp8, param_names, flash_attn_bwd_fp8 = _run_mha_fp8_vs_f16( dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe @@ -2042,6 +2050,7 @@ def test_mha_fp8_vs_f16( os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") fused_attn_fwd_fp8, param_names, fused_attn_bwd_fp8 = _run_mha_fp8_vs_f16( dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe @@ -2051,6 +2060,7 @@ def test_mha_fp8_vs_f16( os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = False") fused_attn_fwd_f16, param_names, fused_attn_bwd_f16 = _run_mha_fp8_vs_f16( dtype, config, False, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe @@ -2290,6 +2300,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FlashAttention)") flash_attn_fwd_fp8, flash_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( dtype, config, True, qkv_layout, is_training, fp8_recipe @@ -2299,6 +2310,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (UnfusedDotProductAttention)") unfused_attn_fwd_fp8, unfused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( dtype, config, True, qkv_layout, is_training, fp8_recipe @@ -2308,6 +2320,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FusedAttention)") fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( dtype, config, True, qkv_layout, is_training, fp8_recipe @@ -2637,6 +2650,8 @@ def _run_custom_mha_fp8(dtype, config, backend): os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + inp = 0.0001 * torch.randint( -100, 100, @@ -2693,6 +2708,8 @@ def _run_ref_mha_f16(dtype, config, backend): os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention": os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + inp = torch.load("qkv.pt").to(device="cuda") inp.requires_grad = True seqlens = torch.full([config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda") diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index fe8f416af4..cdf93d542e 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -453,10 +453,6 @@ def test(): if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() - # F16_arbitrary_seqlen and FP8 are mutually exclusive for a given config (selected by the - # FP8/dtype gate), so a single probe returns the one applicable fused sub-backend. The old - # loop force-set NVTE_FUSED_ATTN_BACKEND per sub-backend, but the refactored backend selection - # no longer reads that env var, so the forcing was inert (and leaked the env var). _attention_backends["backend_selection_requires_update"] = True available_backends, flash_attention_backend, fused_attention_backend = test() if fused_attention_backend in (FusedAttnBackend[name] for name in backends.values()): diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index cf5e0f3475..85435e81e5 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -30,8 +30,7 @@ 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. +// Forward declarations size_t get_max_batch_size(size_t batch_size); size_t get_max_tokens(size_t num_tokens); @@ -97,17 +96,12 @@ void FusedAttnConfig::derive() { FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig cache_cfg = *this; - // Normalize bottom_right_diagonal (the cuDNN diagonal alignment). The impl only turns it into a - // real causal band under `is_causal || is_causal_bottom_right` or a sliding window; otherwise the - // alignment is inert, so canonicalize it (like attn_scale) to false. This keeps the backend - // support probe (which passes a possibly-different brd, e.g. default false) and the real op on a - // single cached graph. + // Normalize bottom_right_diagonal const bool has_window = cache_cfg.window_size_left != -1 || cache_cfg.window_size_right != -1; if (!cache_cfg.is_causal && !cache_cfg.is_causal_bottom_right && !has_window) { cache_cfg.bottom_right_diagonal = false; } else if (cache_cfg.is_causal_bottom_right && cache_cfg.max_seqlen_q == cache_cfg.max_seqlen_kv && !cache_cfg.is_padding) { - // square bottom-right causal collapses to top-left causal (mirrors the impl). cache_cfg.bottom_right_diagonal = false; } diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index ab01de9b91..4469af20bb 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -5,7 +5,7 @@ ************************************************************************/ /*! \file config_and_params.h - * \brief Internal backing objects for fused-attention config and parameter handles. + * \brief Internal objects for fused-attention config and parameter handles. */ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ @@ -76,7 +76,7 @@ struct FusedAttnConfig { // Internal-only fields: never part of attribute serialization, operator<, or the graph cache key. // Filled by derive() or set by caller (i.e. is_forward). Added for convinence purposes and do not - // represent graph properties. + // represent any graph properties. // Direction to build the cuDNN graph for; steers make_cache_key() normalization. bool is_forward = false; diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 159063fd27..56912f5826 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -399,9 +399,9 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); FusedAttnConfig cfg = p.make_config(); - NVTE_Fused_Attn_Backend fused_attention_backend = - nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), - /*message=*/nullptr); + const char *fused_attn_reject_reason = nullptr; + NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( + reinterpret_cast(&cfg), &fused_attn_reject_reason); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { fused_attn_arbitrary_seqlen_fwd(cfg, input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset, @@ -414,7 +414,11 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, p.stream, handle); } else { - NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); + const char *reject_reason = + (fused_attn_reject_reason != nullptr && fused_attn_reject_reason[0] != '\0') + ? fused_attn_reject_reason + : "no cuDNN fused-attention backend supports the requested parameters"; + NVTE_ERROR("Fused attention is not supported for this configuration: ", reject_reason); } } @@ -496,9 +500,9 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); FusedAttnConfig cfg = p.make_config(); - NVTE_Fused_Attn_Backend fused_attention_backend = - nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), - /*message=*/nullptr); + const char *fused_attn_reject_reason = nullptr; + NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( + reinterpret_cast(&cfg), &fused_attn_reject_reason); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; @@ -533,7 +537,11 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, p.stream, handle); } else { - NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); + const char *reject_reason = + (fused_attn_reject_reason != nullptr && fused_attn_reject_reason[0] != '\0') + ? fused_attn_reject_reason + : "no cuDNN fused-attention backend supports the requested parameters"; + NVTE_ERROR("Fused attention is not supported for this configuration: ", reject_reason); } } 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 176773fb51..84c46dcdd1 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 @@ -91,8 +91,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. - // Defined on FusedAttnConfig so make_cache_key() keys the graph on the matching batch - // handling (real batch here, bucketed batch on the legacy path); keep the two in sync. const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; // keep original batch size because cu_seqlens are created with [b+1] shape @@ -153,18 +151,14 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr>; // dropout_offset using CacheType = std::map; - // [SHARED-CACHE] Process-wide graph cache (was `static thread_local`) so a compiled graph - // is reused across threads instead of rebuilt per thread. Safe because cuDNN >= 9.0 allows - // concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe - // execute(). The TE minimum-cuDNN-version bump that formalizes this requirement is a follow-up PR. + // Process-wide graph cache so a compiled graph is reused across threads instead of rebuilt per thread. + // Safe because cuDNN >= 9.0 allows concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe execute(). static CacheType sdpa_f16_fprop_cache; static std::mutex sdpa_f16_fprop_cache_mutex; // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building - // so concurrent first-misses on different keys build in parallel. graph->execute() runs - // unlocked after get_graph() returns; built graphs are shared across threads. + // Lock the map lookup, not the build, so different graphs can build in parallel graph_and_tensors cached_graph{}; bool cache_hit = false; { @@ -443,15 +437,15 @@ void fused_attn_arbitrary_seqlen_fwd_impl( NVTE_CHECK_CUDNN_FE(mha_graph->validate()); NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, - // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. { std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); @@ -704,14 +698,12 @@ void fused_attn_arbitrary_seqlen_bwd_impl( std::shared_ptr>; // dropout_offset using CacheType = std::map; - static CacheType sdpa_f16_bprop_cache; // [SHARED-CACHE] process-wide (was thread_local) - static std::mutex sdpa_f16_bprop_cache_mutex; // [SHARED-CACHE] + static CacheType sdpa_f16_bprop_cache; + static std::mutex sdpa_f16_bprop_cache_mutex; // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building - // so concurrent first-misses on different keys build in parallel. graph->execute() runs - // unlocked after get_graph() returns; built graphs are shared across threads. + // Lock the map lookup, not the build, so different graphs can build in parallel graph_and_tensors cached_graph{}; bool cache_hit = false; { @@ -962,14 +954,14 @@ void fused_attn_arbitrary_seqlen_bwd_impl( NVTE_CHECK_CUDNN_FE(mha_graph->validate()); NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); - // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, - // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. { std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 682af81f55..37082ed9cf 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -126,18 +126,14 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de std::shared_ptr>; // dropout_offset using CacheType = std::map; - // [SHARED-CACHE] Process-wide graph cache (was `static thread_local`) so a compiled graph - // is reused across threads instead of rebuilt per thread. Safe because cuDNN >= 9.0 allows - // concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe - // execute(). The TE minimum-cuDNN-version bump that formalizes this requirement is a follow-up PR. + // Process-wide graph cache so a compiled graph is reused across threads instead of rebuilt per thread. + // Safe because cuDNN >= 9.0 allows concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe execute(). static CacheType sdpa_fp8_fprop_cache; static std::mutex sdpa_fp8_fprop_cache_mutex; // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building - // so concurrent first-misses on different keys build in parallel. graph->execute() runs - // unlocked after get_graph() returns; built graphs are shared across threads. + // Lock the map lookup, not the build, so different graphs can build in parallel graph_and_tensors cached_graph{}; bool cache_hit = false; { @@ -400,13 +396,13 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de NVTE_CHECK_CUDNN_FE(mha_graph->validate()); NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, - // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. { std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); @@ -618,14 +614,12 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr>; // dropout_offset using CacheType = std::map; - static CacheType sdpa_fp8_bprop_cache; // [SHARED-CACHE] process-wide (was thread_local) - static std::mutex sdpa_fp8_bprop_cache_mutex; // [SHARED-CACHE] + static CacheType sdpa_fp8_bprop_cache; + static std::mutex sdpa_fp8_bprop_cache_mutex; // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - // [SHARED-CACHE] Lock only the map lookup; copy the entry out and release before building - // so concurrent first-misses on different keys build in parallel. graph->execute() runs - // unlocked after get_graph() returns; built graphs are shared across threads. + // Lock the map lookup, not the build, so different graphs can build in parallel graph_and_tensors cached_graph{}; bool cache_hit = false; { @@ -1017,14 +1011,14 @@ void fused_attn_fp8_bwd_impl( NVTE_CHECK_CUDNN_FE(mha_graph->validate()); NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); // no-handle overload (handle version is deprecated) - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); // no-handle overload (handle version is deprecated) + NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); + NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); - // [SHARED-CACHE] Lock only for insert. If another thread inserted this key while we built, - // reuse theirs and discard ours so all threads share one graph (rare duplicate build). + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. { std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 0b2dfd6b85..e53ddf92bf 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -8,8 +8,8 @@ * \brief Enums and functions for fused attention. */ -#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_FP8_H_ -#define TRANSFORMER_ENGINE_FUSED_ATTN_FP8_H_ +#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_H_ +#define TRANSFORMER_ENGINE_FUSED_ATTN_H_ #include @@ -466,6 +466,16 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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); +/*! \brief Compute dot product attention with separate Q, K and V. + * + * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_fwd_params()``, + * populate it with ``nvte_set_fused_attn_fwd_params_attribute()`` (or ``FusedAttnFwdParamsWrapper``) setters, and + * destroy it with ``nvte_destroy_fused_attn_fwd_params()``. + * + * \param[in,out] params Opaque fused-attention forward-parameter handle. + */ +void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); + /*! \brief Compute dot product attention with separate Q, K and V. * * Computes: @@ -542,16 +552,6 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); -/*! \brief Compute dot product attention with separate Q, K and V. - * - * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_fwd_params()``, - * populate it with ``nvte_set_fused_attn_fwd_params_attribute()`` (or ``FusedAttnFwdParamsWrapper``) setters, and - * destroy it with ``nvte_destroy_fused_attn_fwd_params()``. - * - * \param[in,out] params Opaque fused-attention forward-parameter handle. - */ -void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); - /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_bwd_params()``, @@ -1139,8 +1139,8 @@ class FusedAttnConfigWrapper { } private: - // Common implementation for every setter: copy the value to a local, forward - // its address and size to the C API, and return *this for chaining. + // Common implementation for every setter: copy the value to a local variable, + // forward its address and size to the C API, and return *this for chaining. template FusedAttnConfigWrapper &set_attr(NVTEFusedAttnConfigAttribute attr, T val) noexcept { nvte_set_fused_attn_config_attribute(cfg_, attr, &val, sizeof(val)); @@ -1287,8 +1287,8 @@ class FusedAttnFwdParamsWrapper { } private: - // Common implementation for every setter: copy the value to a local, forward - // its address and size to the C API, and return *this for chaining. + // Common implementation for every setter: copy the value to a local variable, + // forward its address and size to the C API, and return *this for chaining. template FusedAttnFwdParamsWrapper &set_attr(NVTEFusedAttnFwdParamsAttribute attr, T val) noexcept { nvte_set_fused_attn_fwd_params_attribute(params_, attr, &val, sizeof(val)); @@ -1447,8 +1447,8 @@ class FusedAttnBwdParamsWrapper { } private: - // Common implementation for every setter: copy the value to a local, forward - // its address and size to the C API, and return *this for chaining. + // Common implementation for every setter: copy the value to a local variable, + // forward its address and size to the C API, and return *this for chaining. template FusedAttnBwdParamsWrapper &set_attr(NVTEFusedAttnBwdParamsAttribute attr, T val) noexcept { nvte_set_fused_attn_bwd_params_attribute(params_, attr, &val, sizeof(val)); diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 744e81d7d8..af1eda478d 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -348,10 +348,6 @@ def is_fused_attn_kernel_available( ): """ To check whether the fused attention kernel is supported. - - For a ``POST_SCALE_BIAS`` config, pass the bias broadcast shape via ``bias_batch``, - ``bias_heads``, ``bias_seqlen_q``, and ``bias_seqlen_kv`` so the backend probe matches the - graph used at execution time. """ window_size_tuple = (-1, -1) if window_size is None else window_size diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index dc8aca409a..318d0c15ab 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -136,11 +136,7 @@ class FusedAttnHelper: bias_seqlen_kv: Optional[int] = None def is_fused_attn_kernel_available(self): - """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. - """ + """Check if there is available fused attention kernel""" backend, _ = self.get_fused_attn_backend() return backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 276e476aeb..81cde54adb 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -800,8 +800,6 @@ 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 - # Thread the POST_SCALE_BIAS broadcast shape through so this pre-check probes the same - # cuDNN graph as the primitive does at trace time (see FusedAttnFwdPrimitive.abstract). bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None if attn_bias_type == AttnBiasType.POST_SCALE_BIAS: *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias.shape @@ -846,8 +844,7 @@ def __call__( reason = fused_attn_reject_reason or "(no diagnostic message available)" warnings.warn( "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" + f" support this config. Reason: {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 609e4ef550..631f65d55d 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 @@ -63,14 +63,7 @@ # Setup Attention Logging attn_log.setup_logging() -# Global vars for available attention backends and ALiBi cache. -# -# `_attention_backends` holds the most-recently-selected backend result plus the -# `backend_selection_requires_update` flag. The flag is the public invalidation signal: external -# callers (e.g. the test suite) set it to True to force a full re-selection, typically because they -# changed an NVTE_* environment toggle that is not captured by AttentionParams. This dict is kept -# for backward compatibility (its shape and the flag are part of the de-facto public API); the -# actual multi-entry caching lives in `_attention_backend_cache` below. +# Global vars for available attention backends and ALiBi cache _attention_backends = { "attention_params": None, "use_flash_attention": None, @@ -81,87 +74,6 @@ "backend_selection_requires_update": False, } -# LRU cache of backend-selection results, so that alternating between a handful of configs in the -# same run does not repay get_attention_backend() on every switch (the previous single-slot cache -# thrashed whenever two or more configs interleaved). AttentionParams is unhashable -- it holds -# dicts/lists/tensors and its custom __eq__ disables __hash__ -- so we cannot use it as a dict key. -# Instead we keep an insertion-ordered list of {"attention_params", "env_key", } and -# linear-scan. Capacity is small (10), so the scan is negligible next to a real selection. -# -# The cache identity is (env_key, attention_params). env_key captures the NVTE_* environment toggles -# that get_attention_backend() reads at call time but that AttentionParams does not encode. Including -# it means flipping any such toggle naturally misses and re-selects, so callers do NOT need to -# manually invalidate after changing the environment. Setting -# _attention_backends["backend_selection_requires_update"] = True still hard-clears the whole cache -# for anyone who wants to start completely afresh (e.g. after changing GPU/arch mid-process). -_ATTENTION_BACKEND_RESULT_KEYS = ( - "use_flash_attention", - "flash_attention_backend", - "use_fused_attention", - "fused_attention_backend", - "use_unfused_attention", -) -_ATTENTION_BACKEND_CACHE_MAXSIZE = 10 -_attention_backend_cache = [] - -# Explicit allow-list of the NVTE_* toggles that steer get_attention_backend(). We deliberately do -# NOT snapshot the whole NVTE_* namespace: unrelated toggles (determinism, non-attention modules like -# Linear, debug/logging, etc.) would otherwise needlessly invalidate cached selections. -# -# IMPORTANT: keep this in sync with the os.getenv(...) reads inside -# dot_product_attention/utils.py::get_attention_backend(). If that function begins consulting a new -# NVTE_* toggle that is not listed here, the cache can return a stale (wrong) backend selection. -_ATTENTION_BACKEND_ENV_VARS = ( - "NVTE_FLASH_ATTN", - "NVTE_FLASH_ATTN_V2", - "NVTE_FLASH_ATTN_V3", - "NVTE_FLASH_ATTN_V4", - "NVTE_FUSED_ATTN", - "NVTE_UNFUSED_ATTN", - "NVTE_FP8_DPA_BWD", - "NVTE_DPA_FP8CS_O_in_F16", - "NVTE_DPA_FP8_RECIPE", - "NVTE_DPA_FP8_FORMAT", - "NVTE_DPA_FP8DS_AMAX_ALGO", - "NVTE_DPA_FP8DS_AMAX_HISTLEN", - "NVTE_DPA_FP8DS_REDUCE_AMAX", - "NVTE_UnfusedDPA_Emulate_FP8", -) - - -def _attention_env_key(): - """Snapshot of the selection-relevant NVTE_* toggles (see _ATTENTION_BACKEND_ENV_VARS). - - These influence get_attention_backend() but are not captured by AttentionParams, so they must be - part of the cache identity to avoid returning a result computed under a different environment. A - value of None means the variable is unset (i.e. get_attention_backend() would use its default). - """ - return tuple(os.environ.get(name) for name in _ATTENTION_BACKEND_ENV_VARS) - - -def _attention_backend_cache_lookup(attention_params, env_key): - """Return the cached result dict matching ``(env_key, attention_params)`` (promoted to MRU).""" - for i, entry in enumerate(_attention_backend_cache): - # Compare the cheap env_key tuple before the per-field AttentionParams.__eq__. - if entry["env_key"] == env_key and entry["attention_params"] == attention_params: - if i != len(_attention_backend_cache) - 1: - _attention_backend_cache.append(_attention_backend_cache.pop(i)) - return entry - return None - - -def _attention_backend_cache_store(attention_params, env_key, result): - """Insert/refresh the entry for ``(env_key, attention_params)`` as MRU and evict beyond capacity.""" - for i, entry in enumerate(_attention_backend_cache): - if entry["env_key"] == env_key and entry["attention_params"] == attention_params: - _attention_backend_cache.pop(i) - break - entry = {"attention_params": attention_params, "env_key": env_key, **result} - _attention_backend_cache.append(entry) - while len(_attention_backend_cache) > _ATTENTION_BACKEND_CACHE_MAXSIZE: - _attention_backend_cache.pop(0) - return entry - _alibi_cache = { "_num_heads": None, "_alibi_slopes": None, @@ -1128,13 +1040,12 @@ def forward( Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`, and :attr:`NVTE_UNFUSED_ATTN` to control which DotProductAttention backend to use. - Transformer Engine first filters - backends by support for the runtime environment and input configuration, then applies - a performance-based preference order. On supported pre-Hopper GPUs, FlashAttention is - preferred over FusedAttention and UnfusedDotProductAttention when both optimized - backends are eligible. On Hopper and newer GPUs, including Blackwell, FusedAttention is - preferred over FlashAttention and UnfusedDotProductAttention when both optimized - backends are eligible. + Transformer Engine first filters backends by support for the runtime environment + and input configuration, then applies a performance-based preference order. + On supported pre-Hopper GPUs, FlashAttention is preferred over FusedAttention and + UnfusedDotProductAttention when both optimized backends are eligible. On Hopper and + newer GPUs, including Blackwell, FusedAttention is preferred over FlashAttention and + UnfusedDotProductAttention when both optimized backends are eligible. If FusedAttention is being used, users can also choose to switch to flash-attn's implementation for backward by setting :attr:`NVTE_FUSED_ATTN_USE_FAv2_BWD=1` (default: 0), because of the performance differences between various versions of @@ -1723,17 +1634,13 @@ def forward( use_fused_attention = False use_unfused_attention = True else: - # A forced update hard-clears the entire cache. This is optional now that the - # cache identity includes the NVTE_* environment (so env changes miss on their own); - # it remains as an explicit "start completely afresh" hook (e.g. after changing - # GPU/arch mid-process) for callers who want it. + if ( + _attention_backends["attention_params"] is None + or attention_params != _attention_backends["attention_params"] + ): + _attention_backends["attention_params"] = attention_params + _attention_backends["backend_selection_requires_update"] = True if _attention_backends["backend_selection_requires_update"]: - _attention_backend_cache.clear() - _attention_backends["backend_selection_requires_update"] = False - - env_key = _attention_env_key() - cached = _attention_backend_cache_lookup(attention_params, env_key) - if cached is None: ( use_flash_attention, flash_attention_backend, @@ -1742,17 +1649,14 @@ def forward( use_unfused_attention, _, ) = dpa_utils.get_attention_backend(attention_params) - cached = _attention_backend_cache_store( - attention_params, - env_key, - { - "use_flash_attention": use_flash_attention, - "flash_attention_backend": flash_attention_backend, - "use_fused_attention": use_fused_attention, - "fused_attention_backend": fused_attention_backend, - "use_unfused_attention": use_unfused_attention, - }, - ) + # Set global _attention_backends var using return value + # from get_attention_backend() + _attention_backends["use_flash_attention"] = use_flash_attention + _attention_backends["flash_attention_backend"] = flash_attention_backend + _attention_backends["use_fused_attention"] = use_fused_attention + _attention_backends["fused_attention_backend"] = fused_attention_backend + _attention_backends["use_unfused_attention"] = use_unfused_attention + _attention_backends["backend_selection_requires_update"] = False if use_flash_attention: self.logger.info( "Running with FlashAttention backend (version %s)", @@ -1766,17 +1670,11 @@ def forward( elif use_unfused_attention: self.logger.info("Running with UnfusedDotProductAttention backend") else: - use_flash_attention = cached["use_flash_attention"] - flash_attention_backend = cached["flash_attention_backend"] - use_fused_attention = cached["use_fused_attention"] - fused_attention_backend = cached["fused_attention_backend"] - use_unfused_attention = cached["use_unfused_attention"] - - # Mirror the active selection into the legacy single-slot dict so its public shape - # (and any external readers) keep working as before. - _attention_backends["attention_params"] = attention_params - for _key in _ATTENTION_BACKEND_RESULT_KEYS: - _attention_backends[_key] = cached[_key] + use_flash_attention = _attention_backends["use_flash_attention"] + flash_attention_backend = _attention_backends["flash_attention_backend"] + use_fused_attention = _attention_backends["use_fused_attention"] + fused_attention_backend = _attention_backends["fused_attention_backend"] + use_unfused_attention = _attention_backends["use_unfused_attention"] # raise exception if no backend is available if sum([use_flash_attention, use_fused_attention, use_unfused_attention]) == 0: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 694038752d..53dcbd49b3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1530,7 +1530,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, ) - + # Context-parallel per-step configs if context_parallel: from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( cp_per_step_configs, @@ -2471,11 +2471,7 @@ class FusedAttnSpec: def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_dtype=None): - """Resolve fused-attention specs, e.g. tensor dtypes, formats, for a given config. - - `nominal_dtype` is the model precision (F16/BF16) of the tensors that stay unquantized in - FP8 attention (O, and dQ/dK/dV under current/mxfp8). It is only consulted when `qkv_dtype` itself is FP8. - """ + """Resolve fused-attention specs, e.g. tensor dtypes, formats, for a given config""" q_format = get_qkv_format(qkv_layout)[1] eff_qkv_layout = qkv_layout # FP16/BF16 if recipe is not None: From f802afc371f740609a2ca87eff23bd82743b811b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:55:38 +0000 Subject: [PATCH 41/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/fused_attn_fp8.h | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 4a408772e3..79b279a833 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -19,23 +19,24 @@ namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd(const fused_attn::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, +void fused_attn_fp8_fwd(const fused_attn::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(const fused_attn::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); -// fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd(const fused_attn::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 explaining why it is not supported. From d65c61717e46813d3144a0b8edc3738290480907 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 03:24:52 -0700 Subject: [PATCH 42/63] guard against pre-scale bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/common/fused_attn/fused_attn.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 56912f5826..f27aac426c 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -279,6 +279,12 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } + // cuDNN does not support pre-scale bias + if (cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) { + set_message(message, "Fused attention does not support pre-scale bias."); + 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 = From 47421b967b414a2904c3f9be086731aac655674a Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 03:25:15 -0700 Subject: [PATCH 43/63] fix score mod Jax tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn_score_mod.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index e965a08665..f854f4b16a 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -537,7 +537,7 @@ def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch): assert captured["attn_bias_type"] is AttnBiasType.NO_BIAS assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD assert captured["softmax_type"] is AttnSoftmaxType.VANILLA_SOFTMAX - assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD + assert captured["kernel_checks"][0][0][4] is QKVLayout.BSHD_BSHD_BSHD def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch): @@ -561,7 +561,7 @@ def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monke assert captured["qkv"][0].shape == (1, 8, 1, 16) assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD assert captured["score_mod"] is _identity_score_mod - assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD + assert captured["kernel_checks"][0][0][4] is QKVLayout.BSHD_BSHD_BSHD def test_multi_head_attention_plumbs_score_mod_to_dot_product_attention(monkeypatch): From b6d04eb7c1a73c5ab9fcd26733c7a15903280e08 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 06:32:36 -0700 Subject: [PATCH 44/63] Mirror PyTorch NVTE_DEBUG logging in JAX fused-attn backend selection, document diagnostic semantics, add reject-path test for pre-scale bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 4 +- tests/jax/test_fused_attn_score_mod.py | 100 ++++++++++++++++++ .../common/fused_attn/fused_attn.cpp | 8 ++ .../jax/cpp_extensions/attention.py | 62 ++++++++++- transformer_engine/jax/flax/transformer.py | 5 +- 5 files changed, 174 insertions(+), 5 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index bf32df8971..2685c6deb2 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -403,13 +403,13 @@ Debugging and Profiling :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable debug mode. When set to ``1``, enables verbose debug output and additional checks in attention operations. + :Description: Enable debug mode. When set to ``1``, enables verbose debug output and additional checks in attention operations. Acts as the master switch for the attention backend-selection diagnostics gated by :envvar:`NVTE_DEBUG_LEVEL` (applies to both the PyTorch and JAX fused-attention backends). .. envvar:: NVTE_DEBUG_LEVEL :Type: ``int`` (0, 1, or 2) :Default: ``0`` - :Description: Debug verbosity level. Higher values enable more verbose debug output. Only effective when :envvar:`NVTE_DEBUG` is set to ``1``. + :Description: Debug verbosity level. Higher values enable more verbose debug output. Only effective when :envvar:`NVTE_DEBUG` is set to ``1``. For fused attention, ``1`` logs the outcome (the selected backend, or that none is available) and ``2`` additionally logs the resolved config and the reason fused attention was rejected. This matches the PyTorch attention logging (level 1 = outcome, level 2 = why). .. envvar:: NVTE_PRINT_LAYER_NUMBER diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index f854f4b16a..e0752ff21d 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -17,6 +17,7 @@ QKVLayout, ) from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config +from transformer_engine.jax.cpp_extensions.attention import FusedAttnHelper from transformer_engine.jax.flax import transformer as flax_transformer from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend from test_fused_attn import FusedAttnRunner, SeqDescFormat @@ -762,3 +763,102 @@ def test_fused_attn_score_mod_softcap_with_bprop(): ScoreModFusedAttnRunner.require_cudnn_frontend() runner = ScoreModFusedAttnRunner.softcap(1, 16, 2, 64, jnp.float16) runner.test_backward() + + +def _rejected_fused_attn_helper(): + """A FusedAttnHelper config the backend always rejects (pre-scale bias). + + pre-scale bias is rejected by a TE-side guard in nvte_get_fused_attn_backend_v2 + before any cuDNN probe, so this is hardware- and cuDNN-version-independent. + """ + return FusedAttnHelper( + is_training=True, + batch_size=1, + q_dtype=jnp.float16, + kv_dtype=jnp.float16, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + attn_bias_type=AttnBiasType.PRE_SCALE_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=2, + kv_num_heads=2, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=128, + head_dim_v=128, + window_size=(-1, -1), + bottom_right_diagonal=False, + ) + + +def test_fused_attn_backend_reject_message_bubbles_up(capsys): + """The C++ backend reject reason propagates to Python and reads sensibly. + + Run with `-s` to eyeball the message: + pytest tests/jax/test_fused_attn_score_mod.py -s \ + -k test_fused_attn_backend_reject_message_bubbles_up + """ + helper = _rejected_fused_attn_helper() + backend, message = helper.get_fused_attn_backend() + + with capsys.disabled(): + print(f"\nbackend={backend!r}\nreject_message={message!r}\n") + + assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message, "expected a non-empty diagnostic when the config is rejected" + assert "pre-scale bias" in message.lower() + + +def _cudnn_rejected_fused_attn_helper(): + """A FusedAttnHelper config rejected by the cuDNN is_supported_* probe (not a TE-side guard). + + head_dim=129 is not a multiple of 8, which cuDNN's fused SDPA rejects on every architecture. + This config passes all TE-side guards, so the diagnostic originates from the cuDNN graph-build + exception surfaced via is_supported_f16_fwd(). The exact wording is cuDNN-version dependent. + """ + return FusedAttnHelper( + is_training=True, + batch_size=1, + q_dtype=jnp.float16, + kv_dtype=jnp.float16, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=2, + kv_num_heads=2, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=129, + head_dim_v=129, + window_size=(-1, -1), + bottom_right_diagonal=False, + ) + + +@pytest.mark.skipif( + get_device_compute_capability(0) < 80, + reason="cuDNN fused attention (and its reject path) requires sm80+", +) +def test_fused_attn_backend_cudnn_reject_message_bubbles_up(capsys): + """A cuDNN-level rejection (is_supported_* probe) propagates to Python and reads sensibly. + + Unlike the pre-scale-bias test, this exercises the cuDNN graph-build failure path rather than a + TE-side guard, so the message text is arch/cuDNN-version dependent and only loosely asserted. + + Run with `-s` to eyeball the message: + pytest tests/jax/test_fused_attn_score_mod.py -s \ + -k test_fused_attn_backend_cudnn_reject_message_bubbles_up + """ + helper = _cudnn_rejected_fused_attn_helper() + backend, message = helper.get_fused_attn_backend() + + with capsys.disabled(): + print(f"\nbackend={backend!r}\nreject_message={message!r}\n") + + assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message, "expected a non-empty diagnostic when the config is rejected" + # Not a TE-side guard message -- confirms the reason came from the cuDNN probe path. + assert "pre-scale bias" not in message.lower() diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index f27aac426c..2f5575a8dd 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -234,6 +234,7 @@ 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`. +// NOTE: this overwrites (does not append to) the buffer, so only the last-set reason survives. void set_message(const char **message, std::string reason) { fused_attn_backend_message_buffer = std::move(reason); if (message != nullptr) { @@ -244,6 +245,13 @@ void set_message(const char **message, std::string reason) { } // namespace // select a backend for fused attention +// +// Diagnostic (`message`) semantics: FIRST-FAILURE, not accumulative. The checks below are a +// linear chain of `if (bad) { set_message(...); return NVTE_No_Backend; }` guards, so on +// rejection `message` holds only the first failing reason; later guards never run. Likewise the +// fwd probe short-circuits the bwd probe, and the cuDNN-side is_supported_* probes themselves +// report only the first exception they hit. If a config violates several constraints at once, +// callers see just one of them (fix it and re-run to surface the next). NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, const char **message) { using namespace transformer_engine; diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 318d0c15ab..81ad280b16 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. """JAX/TE custom ops for attention""" +import logging import operator import os import warnings @@ -62,6 +63,42 @@ ] +# NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0 +_NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0")) +# NVTE_DEBUG_LEVEL = 0/1/2 # enables increasingly verbose debug messages, default = 0 +_NVTE_DEBUG_LEVEL = int(os.getenv("NVTE_DEBUG_LEVEL", "0")) + + +class AttentionLogging: + """Manage logging for the JAX attention module. + + Mirrors the PyTorch attention logging so that ``NVTE_DEBUG=1`` combined with + ``NVTE_DEBUG_LEVEL=1/2`` surfaces fused-attention backend-selection diagnostics + (e.g. why a configuration was rejected by cuDNN). + """ + + _log_level = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL + _formatter = logging.Formatter("[%(levelname)-8s | %(name)-19s]: %(message)s") + _stream_handler = logging.StreamHandler() + fa_logger = logging.getLogger(__name__) + _is_logging_setup = False + + @staticmethod + def setup_logging(): + """Set up log levels, logger and handlers (idempotent).""" + if AttentionLogging._is_logging_setup: + return + _log_levels = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG} + AttentionLogging._log_level = _log_levels[ + AttentionLogging._log_level if AttentionLogging._log_level in [0, 1, 2] else 2 + ] + AttentionLogging._stream_handler.setFormatter(AttentionLogging._formatter) + AttentionLogging.fa_logger.setLevel(AttentionLogging._log_level) + if not AttentionLogging.fa_logger.hasHandlers(): + AttentionLogging.fa_logger.addHandler(AttentionLogging._stream_handler) + AttentionLogging._is_logging_setup = True + + @partial( jax.tree_util.register_dataclass, data_fields=[], @@ -145,6 +182,11 @@ def get_fused_attn_backend(self): Returns a ``(backend, message)`` tuple. ``message`` is empty on success, otherwise a diagnostic string explaining why the configuration was rejected. + + When ``NVTE_DEBUG=1``, ``NVTE_DEBUG_LEVEL=1`` logs the outcome (the selected backend, or + that no fused backend is available), and ``NVTE_DEBUG_LEVEL=2`` additionally logs the + resolved config and the reason fused attention was rejected. This mirrors the PyTorch + attention logging (level 1 = outcome, level 2 = why). """ q_type = jax_dtype_to_te_dtype(self.q_dtype) bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 @@ -153,7 +195,7 @@ def get_fused_attn_backend(self): bias_heads = self.bias_heads bias_seqlen_q = self.bias_seqlen_q bias_seqlen_kv = self.bias_seqlen_kv - return transformer_engine_jax.get_fused_attn_backend( + backend, message = transformer_engine_jax.get_fused_attn_backend( self.is_training, self.batch_size, q_type, @@ -189,6 +231,24 @@ def get_fused_attn_backend(self): bias_seqlen_kv, ) + # Mirror the PyTorch attention logging semantics: + # level 1 (INFO) -> the outcome (which backend was selected, or that none was) + # level 2 (DEBUG) -> the "why": the resolved config and the rejection reason + AttentionLogging.setup_logging() + logger = AttentionLogging.fa_logger + logger.debug("Running fused attention backend selection with config=%s", self) + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + logger.info("No fused attention backend available; falling back to unfused attention.") + logger.debug( + "Reason fused attention was rejected: %s", + message or "(no diagnostic message available)", + ) + else: + logger.info("Selected fused attention backend: %s", backend) + if message: + logger.debug("Fused attention backend diagnostic message: %s", message) + return backend, message + @staticmethod def is_non_deterministic_allowed(): """Check if non-deterministic kernels are allowed""" diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 81cde54adb..1279826865 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -831,7 +831,9 @@ def __call__( bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, ) - fused_attn_backend, fused_attn_reject_reason = fused_attn_helper.get_fused_attn_backend() + # get_fused_attn_backend() logs the rejection reason under NVTE_DEBUG, so the + # warning below only reports the (unique) configuration and points there for details. + fused_attn_backend, _ = fused_attn_helper.get_fused_attn_backend() has_fused_attn_kernel = fused_attn_backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend if score_mod_requested and not has_fused_attn_kernel: raise ValueError( @@ -841,7 +843,6 @@ 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( "Falling back to the unfused attention backend as fused attention does not" f" support this config. Reason: {reason}\n" From bcbc084475000907be69b71063f455bda6333725 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Jul 2026 08:31:10 -0700 Subject: [PATCH 45/63] tidy up on jax side Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 4 +- tests/jax/test_fused_attn_score_mod.py | 100 ------------------ .../common/fused_attn/fused_attn.cpp | 9 +- .../jax/cpp_extensions/attention.py | 23 ++-- transformer_engine/jax/flax/transformer.py | 4 +- 5 files changed, 11 insertions(+), 129 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 2685c6deb2..bf32df8971 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -403,13 +403,13 @@ Debugging and Profiling :Type: ``int`` (0 or 1) :Default: ``0`` - :Description: Enable debug mode. When set to ``1``, enables verbose debug output and additional checks in attention operations. Acts as the master switch for the attention backend-selection diagnostics gated by :envvar:`NVTE_DEBUG_LEVEL` (applies to both the PyTorch and JAX fused-attention backends). + :Description: Enable debug mode. When set to ``1``, enables verbose debug output and additional checks in attention operations. .. envvar:: NVTE_DEBUG_LEVEL :Type: ``int`` (0, 1, or 2) :Default: ``0`` - :Description: Debug verbosity level. Higher values enable more verbose debug output. Only effective when :envvar:`NVTE_DEBUG` is set to ``1``. For fused attention, ``1`` logs the outcome (the selected backend, or that none is available) and ``2`` additionally logs the resolved config and the reason fused attention was rejected. This matches the PyTorch attention logging (level 1 = outcome, level 2 = why). + :Description: Debug verbosity level. Higher values enable more verbose debug output. Only effective when :envvar:`NVTE_DEBUG` is set to ``1``. .. envvar:: NVTE_PRINT_LAYER_NUMBER diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index e0752ff21d..f854f4b16a 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -17,7 +17,6 @@ QKVLayout, ) from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config -from transformer_engine.jax.cpp_extensions.attention import FusedAttnHelper from transformer_engine.jax.flax import transformer as flax_transformer from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend from test_fused_attn import FusedAttnRunner, SeqDescFormat @@ -763,102 +762,3 @@ def test_fused_attn_score_mod_softcap_with_bprop(): ScoreModFusedAttnRunner.require_cudnn_frontend() runner = ScoreModFusedAttnRunner.softcap(1, 16, 2, 64, jnp.float16) runner.test_backward() - - -def _rejected_fused_attn_helper(): - """A FusedAttnHelper config the backend always rejects (pre-scale bias). - - pre-scale bias is rejected by a TE-side guard in nvte_get_fused_attn_backend_v2 - before any cuDNN probe, so this is hardware- and cuDNN-version-independent. - """ - return FusedAttnHelper( - is_training=True, - batch_size=1, - q_dtype=jnp.float16, - kv_dtype=jnp.float16, - qkv_layout=QKVLayout.BSHD_BSHD_BSHD, - attn_bias_type=AttnBiasType.PRE_SCALE_BIAS, - attn_mask_type=AttnMaskType.NO_MASK, - softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - dropout_probability=0.0, - q_num_heads=2, - kv_num_heads=2, - q_max_seqlen=128, - kv_max_seqlen=128, - head_dim_qk=128, - head_dim_v=128, - window_size=(-1, -1), - bottom_right_diagonal=False, - ) - - -def test_fused_attn_backend_reject_message_bubbles_up(capsys): - """The C++ backend reject reason propagates to Python and reads sensibly. - - Run with `-s` to eyeball the message: - pytest tests/jax/test_fused_attn_score_mod.py -s \ - -k test_fused_attn_backend_reject_message_bubbles_up - """ - helper = _rejected_fused_attn_helper() - backend, message = helper.get_fused_attn_backend() - - with capsys.disabled(): - print(f"\nbackend={backend!r}\nreject_message={message!r}\n") - - assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message, "expected a non-empty diagnostic when the config is rejected" - assert "pre-scale bias" in message.lower() - - -def _cudnn_rejected_fused_attn_helper(): - """A FusedAttnHelper config rejected by the cuDNN is_supported_* probe (not a TE-side guard). - - head_dim=129 is not a multiple of 8, which cuDNN's fused SDPA rejects on every architecture. - This config passes all TE-side guards, so the diagnostic originates from the cuDNN graph-build - exception surfaced via is_supported_f16_fwd(). The exact wording is cuDNN-version dependent. - """ - return FusedAttnHelper( - is_training=True, - batch_size=1, - q_dtype=jnp.float16, - kv_dtype=jnp.float16, - qkv_layout=QKVLayout.BSHD_BSHD_BSHD, - attn_bias_type=AttnBiasType.NO_BIAS, - attn_mask_type=AttnMaskType.NO_MASK, - softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - dropout_probability=0.0, - q_num_heads=2, - kv_num_heads=2, - q_max_seqlen=128, - kv_max_seqlen=128, - head_dim_qk=129, - head_dim_v=129, - window_size=(-1, -1), - bottom_right_diagonal=False, - ) - - -@pytest.mark.skipif( - get_device_compute_capability(0) < 80, - reason="cuDNN fused attention (and its reject path) requires sm80+", -) -def test_fused_attn_backend_cudnn_reject_message_bubbles_up(capsys): - """A cuDNN-level rejection (is_supported_* probe) propagates to Python and reads sensibly. - - Unlike the pre-scale-bias test, this exercises the cuDNN graph-build failure path rather than a - TE-side guard, so the message text is arch/cuDNN-version dependent and only loosely asserted. - - Run with `-s` to eyeball the message: - pytest tests/jax/test_fused_attn_score_mod.py -s \ - -k test_fused_attn_backend_cudnn_reject_message_bubbles_up - """ - helper = _cudnn_rejected_fused_attn_helper() - backend, message = helper.get_fused_attn_backend() - - with capsys.disabled(): - print(f"\nbackend={backend!r}\nreject_message={message!r}\n") - - assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message, "expected a non-empty diagnostic when the config is rejected" - # Not a TE-side guard message -- confirms the reason came from the cuDNN probe path. - assert "pre-scale bias" not in message.lower() diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 2f5575a8dd..8074f936e9 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -234,7 +234,6 @@ 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`. -// NOTE: this overwrites (does not append to) the buffer, so only the last-set reason survives. void set_message(const char **message, std::string reason) { fused_attn_backend_message_buffer = std::move(reason); if (message != nullptr) { @@ -245,13 +244,7 @@ void set_message(const char **message, std::string reason) { } // namespace // select a backend for fused attention -// -// Diagnostic (`message`) semantics: FIRST-FAILURE, not accumulative. The checks below are a -// linear chain of `if (bad) { set_message(...); return NVTE_No_Backend; }` guards, so on -// rejection `message` holds only the first failing reason; later guards never run. Likewise the -// fwd probe short-circuits the bwd probe, and the cuDNN-side is_supported_* probes themselves -// report only the first exception they hit. If a config violates several constraints at once, -// callers see just one of them (fix it and re-run to surface the next). +// the diagnostic message is based on the first failure, not cumulative NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, const char **message) { using namespace transformer_engine; diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 81ad280b16..8f01e8aab3 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -70,17 +70,12 @@ class AttentionLogging: - """Manage logging for the JAX attention module. - - Mirrors the PyTorch attention logging so that ``NVTE_DEBUG=1`` combined with - ``NVTE_DEBUG_LEVEL=1/2`` surfaces fused-attention backend-selection diagnostics - (e.g. why a configuration was rejected by cuDNN). - """ + """Logging for the JAX attention module""" _log_level = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL _formatter = logging.Formatter("[%(levelname)-8s | %(name)-19s]: %(message)s") _stream_handler = logging.StreamHandler() - fa_logger = logging.getLogger(__name__) + logger = logging.getLogger(__name__) _is_logging_setup = False @staticmethod @@ -93,9 +88,9 @@ def setup_logging(): AttentionLogging._log_level if AttentionLogging._log_level in [0, 1, 2] else 2 ] AttentionLogging._stream_handler.setFormatter(AttentionLogging._formatter) - AttentionLogging.fa_logger.setLevel(AttentionLogging._log_level) - if not AttentionLogging.fa_logger.hasHandlers(): - AttentionLogging.fa_logger.addHandler(AttentionLogging._stream_handler) + AttentionLogging.logger.setLevel(AttentionLogging._log_level) + if not AttentionLogging.logger.hasHandlers(): + AttentionLogging.logger.addHandler(AttentionLogging._stream_handler) AttentionLogging._is_logging_setup = True @@ -185,8 +180,7 @@ def get_fused_attn_backend(self): When ``NVTE_DEBUG=1``, ``NVTE_DEBUG_LEVEL=1`` logs the outcome (the selected backend, or that no fused backend is available), and ``NVTE_DEBUG_LEVEL=2`` additionally logs the - resolved config and the reason fused attention was rejected. This mirrors the PyTorch - attention logging (level 1 = outcome, level 2 = why). + resolved config and the reason fused attention was rejected. """ q_type = jax_dtype_to_te_dtype(self.q_dtype) bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 @@ -231,11 +225,8 @@ def get_fused_attn_backend(self): bias_seqlen_kv, ) - # Mirror the PyTorch attention logging semantics: - # level 1 (INFO) -> the outcome (which backend was selected, or that none was) - # level 2 (DEBUG) -> the "why": the resolved config and the rejection reason AttentionLogging.setup_logging() - logger = AttentionLogging.fa_logger + logger = AttentionLogging.logger logger.debug("Running fused attention backend selection with config=%s", self) if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: logger.info("No fused attention backend available; falling back to unfused attention.") diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 1279826865..695d5bce48 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -831,8 +831,6 @@ def __call__( bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, ) - # get_fused_attn_backend() logs the rejection reason under NVTE_DEBUG, so the - # warning below only reports the (unique) configuration and points there for details. fused_attn_backend, _ = fused_attn_helper.get_fused_attn_backend() has_fused_attn_kernel = fused_attn_backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend if score_mod_requested and not has_fused_attn_kernel: @@ -845,7 +843,7 @@ def __call__( if enable_fused_attn and not has_fused_attn_kernel: warnings.warn( "Falling back to the unfused attention backend as fused attention does not" - f" support this config. Reason: {reason}\n" + f" support this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed rejection reason.\n" ) dropout_rng = None From 1a8c0878e479970d36e110ebbb9f446db50d6b40 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:58:11 +0000 Subject: [PATCH 46/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/jax/flax/transformer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 695d5bce48..2ba5b001f7 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -842,8 +842,9 @@ def __call__( if enable_fused_attn and not has_fused_attn_kernel: warnings.warn( - "Falling back to the unfused attention backend as fused attention does not" - f" support this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed rejection reason.\n" + f"Falling back to the unfused attention backend as fused attention does not support" + f" this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed" + f" rejection reason.\n" ) dropout_rng = None From 1632f7b8e91228f92982d795ce5e8ca2c44d024d Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 04:37:12 -0700 Subject: [PATCH 47/63] fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/jax/flax/transformer.py | 6 +- .../dot_product_attention/context_parallel.py | 20 ++-- .../attention/dot_product_attention/utils.py | 100 +++++++++--------- 3 files changed, 63 insertions(+), 63 deletions(-) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 2ba5b001f7..9578219230 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -842,9 +842,9 @@ def __call__( if enable_fused_attn and not has_fused_attn_kernel: warnings.warn( - f"Falling back to the unfused attention backend as fused attention does not support" - f" this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed" - f" rejection reason.\n" + "Falling back to the unfused attention backend as fused attention does not support" + " this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed" + " rejection reason.\n" ) dropout_rng = None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 8b8bb1b779..7bec9d8a2a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4942,16 +4942,16 @@ def cp_per_step_configs( window_left, window_right = window_size def config(mask, s_q, s_kv, heads, gqa, bottom_right): - return dict( - attn_mask_type=mask, - max_seqlen_q=s_q, - max_seqlen_kv=s_kv, - num_attn_heads=heads, - num_gqa_groups=gqa, - window_size_left=window_left, - window_size_right=window_right, - bottom_right_diagonal=bottom_right, - ) + return { + "attn_mask_type": mask, + "max_seqlen_q": s_q, + "max_seqlen_kv": s_kv, + "num_attn_heads": heads, + "num_gqa_groups": gqa, + "window_size_left": window_left, + "window_size_right": window_right, + "bottom_right_diagonal": bottom_right, + } if cp_comm_type == "a2a": # split heads across the cp ranks diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 53dcbd49b3..5b320dc38d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1486,50 +1486,50 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = ( fu_core_attention_bias_shape ) - base_fused_attn_kwargs = dict( - is_training=is_training, - deterministic=deterministic, - cuda_graph=cuda_graph, - return_max_logit=return_max_logit, - attn_mask_type=AttnMaskType[attn_mask_type], - bias_type=AttnBiasType[fu_core_attention_bias_type], - window_size_left=window_size[0], - window_size_right=window_size[1], - bottom_right_diagonal=bottom_right_diagonal, - softmax_type=SoftmaxType[softmax_type], - scaling_mode=scaling_mode, - dropout=attention_dropout, - attn_scale=softmax_scale, - qkv_dtype=qkv_type, - o_dtype=o_type, - do_dtype=do_type, - dqkv_dtype=dqkv_type, - qkv_layout=QKVLayout[spec.qkv_layout], - o_format=QKVFormat[o_format], - do_format=QKVFormat[do_format], - dqkv_layout=QKVLayout[dqkv_layout], - qkv_scale_inv_format=QKVFormat[qkv_scale_inv_format], - do_scale_inv_format=QKVFormat[do_scale_inv_format], - batch_size=batch_size, - num_attn_heads=num_heads, - num_gqa_groups=num_gqa_groups, - head_dim_qk=head_dim_qk, - head_dim_v=head_dim_v, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, - num_pages_k=num_pages_k, - num_pages_v=num_pages_v, - page_size_k=page_size_k, - page_size_v=page_size_v, - max_pages_per_seq_k=max_pages_per_seq_k, - max_pages_per_seq_v=max_pages_per_seq_v, - bias_batch_size=bias_batch_size, - bias_num_heads=bias_num_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, - ) + base_fused_attn_kwargs = { + "is_training": is_training, + "deterministic": deterministic, + "cuda_graph": cuda_graph, + "return_max_logit": return_max_logit, + "attn_mask_type": AttnMaskType[attn_mask_type], + "bias_type": AttnBiasType[fu_core_attention_bias_type], + "window_size_left": window_size[0], + "window_size_right": window_size[1], + "bottom_right_diagonal": bottom_right_diagonal, + "softmax_type": SoftmaxType[softmax_type], + "scaling_mode": scaling_mode, + "dropout": attention_dropout, + "attn_scale": softmax_scale, + "qkv_dtype": qkv_type, + "o_dtype": o_type, + "do_dtype": do_type, + "dqkv_dtype": dqkv_type, + "qkv_layout": QKVLayout[spec.qkv_layout], + "o_format": QKVFormat[o_format], + "do_format": QKVFormat[do_format], + "dqkv_layout": QKVLayout[dqkv_layout], + "qkv_scale_inv_format": QKVFormat[qkv_scale_inv_format], + "do_scale_inv_format": QKVFormat[do_scale_inv_format], + "batch_size": batch_size, + "num_attn_heads": num_heads, + "num_gqa_groups": num_gqa_groups, + "head_dim_qk": head_dim_qk, + "head_dim_v": head_dim_v, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_kv": max_seqlen_kv, + "num_tokens_q": num_tokens_q, + "num_tokens_kv": num_tokens_kv, + "num_pages_k": num_pages_k, + "num_pages_v": num_pages_v, + "page_size_k": page_size_k, + "page_size_v": page_size_v, + "max_pages_per_seq_k": max_pages_per_seq_k, + "max_pages_per_seq_v": max_pages_per_seq_v, + "bias_batch_size": bias_batch_size, + "bias_num_heads": bias_num_heads, + "bias_seqlen_q": bias_seqlen_q, + "bias_seqlen_kv": bias_seqlen_kv, + } # Context-parallel per-step configs if context_parallel: from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( @@ -2482,12 +2482,12 @@ def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_d eff_qkv_layout = qkv_layout # MXFP8 fast path else: eff_qkv_layout = "bhsd_bhsd_bhsd" # MXFP8 slow path - layout_kwargs = dict( - qkv_layout=eff_qkv_layout, - o_format=q_format, - do_format=q_format, - dqkv_layout=qkv_layout, - ) + layout_kwargs = { + "qkv_layout": eff_qkv_layout, + "o_format": q_format, + "do_format": q_format, + "dqkv_layout": qkv_layout, + } if qkv_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): ref = TE_DType[nominal_dtype if nominal_dtype is not None else torch.bfloat16] From c0233a10878ff0160f708bf5e901e6aa1b593c3f Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 04:43:10 -0700 Subject: [PATCH 48/63] fix nvte_get_fused_attn_backend shim, docstring, bias/softmax pointers Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 28 ++++++++----- .../include/transformer_engine/fused_attn.h | 39 ++++++++++++------- 2 files changed, 43 insertions(+), 24 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 8074f936e9..bddaafb8d6 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -228,8 +228,8 @@ 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 +// The per-thread storage for the diagnostic string; it's re-used (cleared + re-populated) +// on every call to nvte_get_fused_attn_backend_v2 on the same 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, @@ -243,8 +243,7 @@ void set_message(const char **message, std::string reason) { } // namespace -// select a backend for fused attention -// the diagnostic message is based on the first failure, not cumulative +// select a backend for fused attention; the diagnostic message is based on the first failure, not cumulative. NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, const char **message) { using namespace transformer_engine; @@ -346,10 +345,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } -// 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. +// 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, @@ -379,10 +375,23 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( cfg.is_training = is_training; cfg.return_max_logit = return_max_logit; cfg.deterministic = deterministic; + // fill in missing fields so it doesn't always return NVTE_No_Backend + cfg.batch_size = 1; + cfg.o_format = nvte_get_q_format(qkv_layout); + cfg.do_format = cfg.o_format; + cfg.dqkv_layout = qkv_layout; + if (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { + cfg.bias_batch_size = cfg.batch_size; + cfg.bias_num_heads = num_attn_heads; + cfg.bias_seqlen_q = max_seqlen_q; + cfg.bias_seqlen_kv = max_seqlen_kv; + } + return nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), /*message=*/nullptr); } +// fused attention forward void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; @@ -482,6 +491,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso nvte_fused_attn_fwd_v2(reinterpret_cast(&p)); } +// fused attention backward void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { NVTE_API_CALL(nvte_fused_attn_bwd_v2); using namespace transformer_engine; @@ -515,7 +525,7 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { size_t i = 0; Tensor *output_S = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - Tensor *input_Bias, *input_SoftmaxOffset; + Tensor *input_Bias = nullptr, *input_SoftmaxOffset = nullptr; if ((p.bias_type != NVTE_NO_BIAS) && (p.bias_type != NVTE_ALIBI)) { input_Bias = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); } diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index e53ddf92bf..938fa1747e 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -413,20 +413,20 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, const void *buf, size_t size_in_bytes); -/*! \brief Get fused-attention backend based on input parameters. +/*! \brief Get fused-attention backend based on user configuration. * - * This function runs cuDNN frontend's support surface checks, builds cuDNN graphs, - * and caches them if the build is successful. + * This function passes the user configuration to cuDNN frontend, runs its support checks, + * attempts to build the necessary graphs, and if successful, caches the graphs (if not, returns + * ``NVTE_No_Backend``). * * \param[in] cfg Fused-attention configuration created by * ``nvte_create_fused_attn_config()``. * \param[out] message If cuDNN graphs are built successfully, an empty string; - * if not, a diagnostic message with the reason for rejection. - * Pass NULL to skip diagnostics. - * 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 + * if not, a diagnostic message explaining why there is no support. + * Pass NULL to skip the diagnostics. Note that 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. * * \return Fused-attention backend, ``NVTE_F16_arbitrary_seqlen`` or ``NVTE_FP8``, @@ -458,6 +458,13 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, * \param[in] deterministic Whether determinism is required or not. * * \deprecated This function has been deprecated in favor of nvte_get_fused_attn_backend_v2. + * + * \note nvte_get_fused_attn_backend has a narrower signature than nvte_get_fused_attn_backend_v2, + * and it fills the fields that it cannot express with default values. For example, it sets + * batch_size = 1, derives output/gradient formats from qkv_layout, assumes a standard + * bias shape [b, h, sq, skv] for NVTE_POST_SCALE_BIAS, uses delayed scaling for all FP8, + * and does not support paged-KV attention. Users who need more precise control should + * use nvte_get_fused_attn_backend_v2 directly. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, @@ -466,10 +473,11 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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); -/*! \brief Compute dot product attention with separate Q, K and V. +/*! \brief Compute dot product attention with Q, K, and V. * - * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_fwd_params()``, - * populate it with ``nvte_set_fused_attn_fwd_params_attribute()`` (or ``FusedAttnFwdParamsWrapper``) setters, and + * All inputs and outputs are carried by the opaque \p params handle. Create it with + * ``nvte_create_fused_attn_fwd_params()``, populate it with + * ``nvte_set_fused_attn_fwd_params_attribute()`` (or ``FusedAttnFwdParamsWrapper``) setters, and * destroy it with ``nvte_destroy_fused_attn_fwd_params()``. * * \param[in,out] params Opaque fused-attention forward-parameter handle. @@ -552,10 +560,11 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); -/*! \brief Compute the backward of the dot product attention with separate Q, K and V. +/*! \brief Compute the backward of the dot product attention with Q, K and V. * - * All inputs and outputs are carried by the opaque \p params handle. Create it with ``nvte_create_fused_attn_bwd_params()``, - * populate it with ``nvte_set_fused_attn_bwd_params_attribute()`` (or ``FusedAttnBwdParamsWrapper``) setters, and + * All inputs and outputs are carried by the opaque \p params handle. Create it with + * ``nvte_create_fused_attn_bwd_params()``, populate it with + * ``nvte_set_fused_attn_bwd_params_attribute()`` (or ``FusedAttnBwdParamsWrapper``) setters, and * destroy it with ``nvte_destroy_fused_attn_bwd_params()``. * * \param[in,out] params Opaque fused-attention backward-parameter handle. From 928dd334df29c31bea5392bfb328e4decf69c6aa Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 05:42:01 -0700 Subject: [PATCH 49/63] add device_id as a key Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.cpp | 3 +++ transformer_engine/common/fused_attn/config_and_params.h | 8 ++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 85435e81e5..ca4214dac3 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -96,6 +96,9 @@ void FusedAttnConfig::derive() { FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig cache_cfg = *this; + // Key the device ID for multi-GPU single-process runs + cache_cfg.device_id = cuda::current_device(); + // Normalize bottom_right_diagonal const bool has_window = cache_cfg.window_size_left != -1 || cache_cfg.window_size_right != -1; if (!cache_cfg.is_causal && !cache_cfg.is_causal_bottom_right && !has_window) { diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 4469af20bb..ebc5b3eb07 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -74,6 +74,10 @@ struct FusedAttnConfig { size_t bias_seqlen_q = 0; size_t bias_seqlen_kv = 0; + // device ID: not part of attribute serialization, but part of operator< and used to + // differentiate graphs built for different devices in multi-GPU single-process runs + int device_id = -1; + // Internal-only fields: never part of attribute serialization, operator<, or the graph cache key. // Filled by derive() or set by caller (i.e. is_forward). Added for convinence purposes and do not // represent any graph properties. @@ -156,7 +160,7 @@ struct FusedAttnConfig { head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_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_batch_size, bias_num_heads, bias_seqlen_q, - bias_seqlen_kv) < + bias_seqlen_kv, device_id) < std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, @@ -167,7 +171,7 @@ struct FusedAttnConfig { rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, 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.bias_num_heads, rhs.bias_seqlen_q, rhs.bias_seqlen_kv, rhs.device_id); } // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields From 7ff80580eaff49a444ea3bc6ce9ec4b8acd719e5 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:18:12 -0700 Subject: [PATCH 50/63] add docstring for FP8 recipes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../attention/dot_product_attention/backends.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 785e438cda..78ae57d849 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1929,7 +1929,19 @@ class FusedAttention(torch.nn.Module): FusedAttnBackend["F16_arbitrary_seqlen"] cuDNN attention for FP16/BF16 with any sequence length. FusedAttnBackend["FP8"] - cuDNN attention for FP8 with any sequence length. + cuDNN attention for FP8 with any sequence length. The supported recipes are as follows. Inputs, + Intermediates, and Outputs are in the format of "tensor: quantizer", and are used by function calls, + tex.fused_attn_fwd and tex.fused_attn_bwd. + + Direction Inputs Intermediates Outputs + DelayedScaling (DS) forward Q/K/V: DS S: DS O: DS + backward Q/K/V/O (from forward), dO: DS dP: DS dQ/dK/dV: DS + Float8CurrentScaling (CS) forward Q/K/V: CS S: DS O: F16 + backward Q/K/V (from forward), dO: CS, + O: F16 (or CS if NVTE_DPA_FP8CS_O_in_F16=0) dP: DS dQ/dK/dV: F16 + MXFP8BlockScaling (MXFP8) forward Q/K row, V col: MXFP8 S: None O: F16 + backward Q/K row+col, V row: MXFP8, + O/dO: F16, dO row+col: MXFP8 dP: None dQ/dK/dV: F16 """ def __init__( From 1db50d9c5559a5408a49efc1341e9859a49e1cfd Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 07:19:20 -0700 Subject: [PATCH 51/63] fix doc/ipynb Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/examples/attention/attention.ipynb | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 6c868518ec..ee6c553bbb 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -175,6 +175,7 @@ "outputs": [ { "output_type": "stream", + "name": "stdout", "text": [ "Device 0: NVIDIA H100 80GB HBM3 GPU, sm90 compute capability, 79.1GB memory\n", "Running test_0 with cuDNN attention and flash-attention...\n", @@ -273,6 +274,7 @@ "outputs": [ { "output_type": "stream", + "name": "stdout", "text": [ "\n", "Run cuDNN attention...\n", @@ -305,6 +307,7 @@ "outputs": [ { "output_type": "stream", + "name": "stdout", "text": [ "\n", "Run cuDNN attention...\n", @@ -508,6 +511,7 @@ "outputs": [ { "output_type": "stream", + "name": "stdout", "text": [ "Run with post_scale_bias:\n", "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", From dca9585458e640aed1143f4a201dcd6745b5c0a3 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:15:06 -0700 Subject: [PATCH 52/63] avoid duplicate checks for fused backend and force to 0 for bias shape defaults Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_distributed_fused_attn.py | 40 ------------------- .../jax/cpp_extensions/attention.py | 8 ++-- 2 files changed, 4 insertions(+), 44 deletions(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index a03f5ad9c2..6657962e93 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -81,26 +81,6 @@ def impl_test_self_attn( is_training = True batch, seqlen, num_head, hidden = data_shape - if not is_fused_attn_kernel_available( - is_training, - batch, - dtype, - dtype, - QKVLayout.BS3HD, - attn_bias_type, - attn_mask_type, - softmax_type, - dropout_prob, - num_head, - num_head, - seqlen, - seqlen, - hidden, - hidden, - None, # no window - ): - pytest.skip("No FusedAttn backend found") - col_ref = self.generate_collectives_count_ref( mesh_shape, mesh_axes, @@ -234,26 +214,6 @@ def test_cross_attn( batch, seqlen, num_head, hidden = data_shape - if not is_fused_attn_kernel_available( - is_training, - batch, - dtype, - dtype, - QKVLayout.BSHD_BS2HD, - attn_bias_type, - attn_mask_type, - softmax_type, - dropout_prob, - num_head, - num_head, - seqlen, - seqlen, - hidden, - hidden, - None, # no window - ): - pytest.skip("No FusedAttn backend found") - col_ref = self.generate_collectives_count_ref() runner = FusedAttnRunner( batch, diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 8f01e8aab3..f8cd1308cb 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -185,10 +185,10 @@ def get_fused_attn_backend(self): q_type = jax_dtype_to_te_dtype(self.q_dtype) bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - bias_batch = self.bias_batch - bias_heads = self.bias_heads - bias_seqlen_q = self.bias_seqlen_q - bias_seqlen_kv = self.bias_seqlen_kv + bias_batch = self.bias_batch or 0 + bias_heads = self.bias_heads or 0 + bias_seqlen_q = self.bias_seqlen_q or 0 + bias_seqlen_kv = self.bias_seqlen_kv or 0 backend, message = transformer_engine_jax.get_fused_attn_backend( self.is_training, self.batch_size, From 7585a1bd6113f11e70fc8e7f0dcad16b94b0c191 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:31:12 -0700 Subject: [PATCH 53/63] reduce ipynb diffs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/examples/attention/attention.ipynb | 1230 +++++++++++------------ 1 file changed, 615 insertions(+), 615 deletions(-) diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index ee6c553bbb..4ffa804401 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -1,622 +1,622 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Attention Is All You Need!\n", - "\n", - "The core idea behind Transformer models is the attention mechanism [[1]](https://arxiv.org/abs/1706.03762). It identifies the correlation between words, selects the most important parts of the sentence to focus on, and captures meaningful patterns and dependencies in the data. Figure 1 shows a typical attention mechanism, where pre-softmax operations can be a combination of scaling, bias and masking while the post-softmax operation is often just dropout.\n", - "\n", - "
    \n", - "\n", - "
    Figure 1: Dot product attention.
    \n", - "
    \n", - "\n", - "[Transformer Engine](https://github.com/NVIDIA/TransformerEngine.git) supports the calculation of dot product attention in two frameworks, [PyTorch](https://github.com/pytorch/pytorch) and [JAX](https://github.com/google/jax). The API for each framework is\n", - "\n", - "- [transformer_engine.pytorch.DotProductAttention](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention)\n", - "- [transformer_engine.jax.flax.DotProductAttention](../../api/jax.rst#transformer_engine.jax.flax.DotProductAttention)" - ], - "id": "040f466a" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 1. Attention Backends\n", - "\n", - "Transformer Engine provides multiple attention backends for each supported framework. The framework-native backends provide a robust baseline, while the fused, GPU-optimized implementations offer more performance. For example, the flash-attention and cuDNN attention backends in PyTorch. The framework-native backends are often named with \"unfused\", while the more optimized backends are \"fused\" or \"flash\".\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    FrameworkBackend (Module Name)Module Location
    PyTorchcuDNN attention (`FusedAttention`) [transformer_engine.pytorch.attention](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py)
    flash-attention (`FlashAttention`)
    \n", - " PyTorch-native attention (`UnfusedDotProductAttention`)\n", - "
    JAXcuDNN attention (`_FusedDotProductAttention`)[transformer_engine.jax.flax.transformer](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/flax/transformer.py)
    JAX-native attention (`_UnfusedDotProductAttention`)
    " - ], - "id": "89a7d849" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 1.1 Flash vs. Non-Flash\n", - "\n", - "The attention calculation has quadratic computational and memory complexities to the sequence length. Its runtime and memory requirements quadruple, when the sequence length doubles. This presents a significant challenge to scale Transformer models up for longer contexts, in order to achieve higher model quality.\n", - "\n", - "Compared to the standard, non-flash algorithm, the flash algorithm [[2]](https://arxiv.org/abs/2205.14135) was proposed to reduce the memory scaling to linear and improve the computational efficiency through optimized memory accesses. It employs the following two distinctive techniques.\n", - "\n", - "- **Tiling:** The non-flash algorithm tries to process the query, key, value tensors in one single step, requiring large amounts of global memory and incurring high volumes of reads/writes between global memory and shared memory. The flash algorithm decomposes the input into several tiles, based on the available shared memory and register size, and it computes the softmax one tile at a time.\n", - "\n", - "- **Recomputation:** The non-flash algorithm stores the softmax matrix (quadratic to sequence length) to global memory for the backward pass, while the flash algorithm only saves the softmax normalization factors (linear to sequence length). This reduces the amount of memory required as well as the bandwidth utilization between global memory and shared memory. Even though there is extra computation incurred in order to recalculate the attention in the backward pass, the bandwidth savings still provide significant improvement in efficiency.\n", - "\n", - "
    \n", - "Note: \n", - " \n", - "Transformer Engine's flash-attention backend, available in PyTorch, and cuDNN attention backend (sub-backends 1 and 2), available in PyTorch and JAX, are both based on the flash algorithm.\n", - "
    \n" - ], - "id": "c90a2573" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 1.2 flash-attention\n", - "\n", - "The flash-attention backend, available only in PyTorch, is a module wrapped around the public `flash-attn` package [[3]](https://github.com/Dao-AILab/flash-attention). \n", - "\n", - "The flash-attention backend supports `flash-attn`'s features as well as a few extra functionalities to facilitate the use of `flash-attn`, such as converting the `attention_mask` to cumulative sequence lengths `cu_seqlens` for `padding` mask use cases. Please see `transformer_engine.pytorch.attention.FlashAttention` for details.\n", - "\n", - "The `flash-attn` dependency is regularly updated in Transformer Engine. As of v2.0, Transformer Engine supports `flash-attn` 2.0.6+ (see [setup.py](https://github.com/NVIDIA/TransformerEngine/blob/main/setup.py)).\n", - "\n", - "To understand `flash-attn`'s performance, please refer to their benchmarks [here](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#performance).\n", - "\n", - "### 1.3 cuDNN Attention\n", - "\n", - "The cuDNN attention backend, available in PyTorch and JAX, offers another high-performance solution to the attention calculation. It requires [cuDNN](https://developer.nvidia.com/cudnn) to run, and has several sub-backends to support the different precisions and sequence lengths.\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    Sub-BackendAlgorithmPrecisionSequence LengthArchitectureAdditional info
    1FlashBF16/FP16 Any sm80+ [cuDNN](https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#fused-flash-attention-fprop),\n", - " [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention)\n", - "
    2FlashFP8 cuDNN pre-9.0: ≤512 cuDNN pre-9.0: sm90
    cuDNN 9.0+: Any cuDNN 9.0+: sm90+ cuDNN 9.0+: [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention-fp8)\n", - "
    \n", - "\n", - "The cuDNN attention backend and flash-attention backend have several notable differences. As of Transformer Engine 2.0, cuDNN 9.3 and `flash-attn` 2.4.2,\n", - "\n", - "- flash-attention only supports the PyTorch framework while cuDNN attention supports PyTorch and JAX.\n", - "- flash-attention supports BF16, FP16 precisions while cuDNN attention also supports FP8 (through its sub-backend 2).\n", - "- flash-attention supports `bshd`, `thd` input formats, without any transposes, and `sbhd` format, with transposes, while cuDNN attention supports all three formats without transposes (see Section 3.1 for more details).\n", - "- flash-attention does not support `post_scale_bias`, and cuDNN attention does.\n", - "- flash-attention supports KV-caching and paged attention, and cuDNN attention does not.\n", - "- flash-attention uses bottom right diagonal for `causal` mask in cross attention (see [change log](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#21-change-behavior-of-causal-flag)), and cuDNN attention supports both top left and bottom right.\n", - "- **Sliding window attention (SWA):** flash-attention has SWA(left, right) support for all mask types except top-left causal masks, with or without dropout, and without bias. cuDNN attention supports SWA(left, 0) starting from 9.2 and SWA(left, right) starting from 9.6, without dropout, and with `bias_type=\"no_bias\"`.\n", - "- flash-attention outperforms cuDNN attention on Ampere architectures, and cuDNN attention has 20-50% advantages on Hopper architectures, based on our benchmarks for a number of commonly-used model configurations.\n", - "\n", - "To compare cuDNN attention and flash-attention, users can modify the `model_configs` dictionary in [benchmarks/attention/benchmark_attention.py](https://github.com/NVIDIA/TransformerEngine/blob/main/benchmarks/attention/benchmark_attention.py) to collect performance numbers. The script runs each entry in `model_configs` for `num_iters` times, each time with one forward pass and one backward pass. Both backends are tried, and if one backend does not have support for the specific user input, the runtimes and speedups in the final table would be 0." - ], - "id": "b5ce567d" - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "model_configs = {\n", - " # test: b, h, hg, d, sq, skv, p, mask, bias\n", - " \"test_0\": ModelConfig(2, 16, 16, 64, 512, 512, 0.0, \"no_mask\", \"no_bias\"), # short seq\n", - " \"test_1\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"no_bias\"), # longer seq, mask\n", - " \"test_2\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"post_scale_bias\"), # bias\n", - " \"test_3\": ModelConfig(2, 32, 4, 128, 8192, 8192, 0.0, \"causal\", \"no_bias\"), # GQA\n", - "}" - ], - "execution_count": null, - "outputs": [], - "id": "c5b8e3d7" - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "!cd ../../../benchmarks/attention/ && python benchmark_attention.py" - ], - "execution_count": null, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Device 0: NVIDIA H100 80GB HBM3 GPU, sm90 compute capability, 79.1GB memory\n", - "Running test_0 with cuDNN attention and flash-attention...\n", - "Running test_1 with cuDNN attention and flash-attention...\n", - "Running test_2 with cuDNN attention...\n", - "Running test_3 with cuDNN attention and flash-attention...\n", - "\n", - " cuDNN fwd+bwd (ms) flash-attn fwd+bwd (ms) cuDNN vs flash speedup\n", - "test_0 0.0340 0.0468 1.3786\n", - "test_1 0.3664 0.5850 1.5968\n", - "test_2 0.9332 0.0000 0.0000\n", - "test_3 7.4875 11.8879 1.5877\n" - ] - } - ], - "id": "50852cb5" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 2. Backend Selection\n", - "\n", - "Given the various attention backends, Transformer Engine first determines which backends are eligible for the provided inputs and runtime environment, then applies a preference order among the eligible backends. Eligibility is affected by user environment variables, GPU architecture, installed `flash-attn` and cuDNN versions, data type and FP8 recipe, QKV layout, training or inference mode, dropout, and other attention features.\n", - "\n", - "In PyTorch, the candidates are FlashAttention (`flash-attn` v2, v3, or v4), FusedAttention (cuDNN sub-backends), and UnfusedDotProductAttention. Users can disable whole backend families with `NVTE_FLASH_ATTN`, `NVTE_FUSED_ATTN`, or `NVTE_UNFUSED_ATTN`. In JAX, Transformer Engine checks whether a cuDNN fused-attention kernel is available when `NVTE_FUSED_ATTN=1`; otherwise it falls back to the JAX-native implementation.\n", - "\n", - "At a high level, the architecture-specific PyTorch selection order is:\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    FrameworkSelection Order
    PyTorchsm8x (Ampere/Ada): flash-attention > cuDNN attention > PyTorch-native attention
    sm90 (Hopper): cuDNN attention > flash-attention > PyTorch-native attention
    sm100/sm120 (Blackwell): cuDNN attention > flash-attention > PyTorch-native attention
    cuDNN attention: BF16/FP16 uses sub-backend 1 when eligible; FP8 uses sub-backend 2 when enabled and eligible
    JAXcuDNN attention > JAX-native attention
    \n", - "\n", - "Within FlashAttention, TE uses the installed implementation that is supported for the architecture and input. FlashAttention 3 is Hopper-only (`sm90`). FlashAttention 4 supports `sm80`, `sm90`, `sm100`, and `sm120`; on Hopper, TE prefers FlashAttention 3 over FlashAttention 4 when both are installed and eligible. On Blackwell, FlashAttention 4 is the Blackwell-specific flash-attention path when installed and eligible, while FlashAttention 2 can still be eligible depending on the installed version and input configuration.\n", - "\n", - "Within cuDNN FusedAttention, TE asks the fused-attention helper which sub-backend is eligible. Sub-backend 1 is the BF16/FP16 flash-based path when available; sub-backend 2 is the FP8 path when FP8 DPA is enabled and the architecture, cuDNN version, and input configuration support it. Hopper supports eligible FP8 DPA through cuDNN sub-backend 2. In the current PyTorch selector, eligible FP8 DPA on Blackwell is an `sm100` path and is disabled on `sm120`.\n", - "\n", - "When all optimized backends are disabled or ineligible, TE falls back to UnfusedDotProductAttention if it is enabled. If no backend is eligible, backend selection returns no backend and the caller raises an error. As we monitor the performance of different backends, the selection logic may change." - ], - "id": "9a615119" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 2.1 Debug Information\n", - "\n", - "To find out which backend is being used during runtime, we have the following two debugging flags. Logging is done by using the `logging` package.\n", - "```\n", - "NVTE_DEBUG = 0/1 # disables/enables debugging\n", - "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n", - "```\n", - "
    \n", - "Note:\n", - " \n", - "These flags are supported in PyTorch only as of Transformer Engine 2.0. JAX support is expected to be added in the future.\n", - "
    " - ], - "id": "e6c0f3f0" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The example script [example_attention.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/example_attention.py) runs a very basic model with two attention backends, cuDNN attention and flash-attention. Here `NVTE_DEBUG_LEVEL=1` allows us to find out which backend/sub-backend is used in runtime." - ], - "id": "16660323" - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python example_attention.py" - ], - "execution_count": null, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "\n", - "Run cuDNN attention...\n", - "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", - "\n", - "Run flash-attention...\n", - "[INFO | DotProductAttention]: Running with FlashAttention backend\n", - "\n", - "Test passed.\n" - ] - } - ], - "id": "906b8cf1" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "`NVTE_DEBUG_LEVEL=2` allows us to find out more about the backend selection logic. Users are encouraged to double check the `config` and provide it to the Transformer Engine team if they would like to file a bug. " - ], - "id": "8ca99461" - }, - { - "cell_type": "code", - "metadata": {}, - "source": [ - "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py" - ], - "execution_count": null, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "\n", - "Run cuDNN attention...\n", - "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", - "[DEBUG | DotProductAttention]: Disabling FlashAttention due to NVTE_FLASH_ATTN=0\n", - "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=False, FusedAttention=True (sub-backend 1), UnfusedDotProductAttention=True}\n", - "[DEBUG | DotProductAttention]: Selected backend = FusedAttention (sub-backend 1)\n", - "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", - "\n", - "Run flash-attention...\n", - "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", - "[DEBUG | DotProductAttention]: Disabling FusedAttention due to NVTE_FUSED_ATTN=0\n", - "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=True, FusedAttention=False, UnfusedDotProductAttention=True}\n", - "[DEBUG | DotProductAttention]: Selected backend = FlashAttention\n", - "[INFO | DotProductAttention]: Running with FlashAttention backend\n", - "\n", - "Test passed.\n" - ] - } - ], - "id": "d3637094" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 2.2 User Control\n", - "\n", - "Users usually do not need to worry about the backend selection. However, if there is a convergence or performance issue encountered, Transformer Engine provides a few other environment variables for users to experiment with different backends.\n", - "\n", - "**flash-attention or cuDNN attention:**\n", - "Users can enable/disable the flash-attention backend or cuDNN attention backend via the following two environment variables in PyTorch.\n", - "```\n", - "NVTE_FLASH_ATTN = 0 # disables flash-attention; default = 1\n", - "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n", - "```\n", - "\n", - "```\n", - "
    \n", - "Note\n", - " \n", - "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", - "
    \n", - "\n", - "### 2.3 Example Tests\n", - "\n", - "Our [unit tests](https://github.com/NVIDIA/TransformerEngine/tree/main/tests) demonstrate the use of Transformer Engine dot product attention APIs. Users are encouraged to use them as a template when integrating Transformer Engine to their ML workflows.\n", - "\n", - "For example, in PyTorch, [test_dot_product_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) offers a variety of use cases of `pytorch.DotProductAttention`, from data types, model configs, checkpointing, to QKV layouts." - ], - "id": "611d8fdb" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 3. Backend Support\n", - "\n", - "Transformer Engine supports commonly-used features such as self and cross attention, FP16/BF16 precisions, dropout, and checkpointing. But it also offers a range of other features. As of v2.0, Transformer Engine's attention backends have the following support matrix.\n", - "\n", - "| Attention Backend | Precision | Architecture | Sliding Window Attention | MQA/GQA | Multi-Latent Attention | Context Parallelism | Determinism Possible |\n", - "| :---------------- | :-------- | :----------- | :----------------------- | :------ | :--------------------- | :------------------ | :------------ |\n", - "| cuDNN attention (all frameworks) | BF16, FP16, FP8 (PyTorch only) | sm80+ | Yes (cuDNN 9.2+) | Yes | Yes | Yes (`bshd`,`sbhd`, `thd`) | Yes |\n", - "| flash-attention (PyTorch) | BF16, FP16 | sm80+ | Yes | Yes | Yes | Yes (`bshd`,`thd`) | Yes |\n", - "| Framework-native attention | BF16, FP16, FP32 | Any | No, unless used as a mask | Yes | Yes (PyTorch only) | No | Yes |\n", - "\n", - "Some unit tests are provided to serve as a starting point for integrating such features into users' models. For example,\n", - "- sliding window attention: [test_dpa_swa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", - "- MQA/GQA: [test_te_layer_mqa_gqa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", - "- Multi-Latent Attention: [test_dpa_mla](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", - "- context parallelism: [test_cp_with_fused_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py), [test_cp_with_flash_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py)" - ], - "id": "e60a2a3e" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 3.1 QKV Layout\n", - "\n", - "Transformer Engine supports various layouts of the query `q`, key `k`, value `v` tensors. It has defined 15 QKV layouts, which are grouped into 3 QKV formats and 5 QKV layout groups to help with similar memory/computational operations across different layouts. The mapping relationships of these layouts and groups are,\n", - "\n", - "| `qkv_layout`         | `qkv_layout_group`=`3hd` | `h3d` | `hd_2hd` | `hd_h2d` | `hd_hd_hd` |\n", - "| ----------: | -----------: | -----: | ----------: | ----------: | -------------: |\n", - "| `qkv_format`=`sbhd` | `sb3hd` | `sbh3d` | `sbhd_sb2hd` | `sbhd_sbh2d` | `sbhd_sbhd_sbhd` |\n", - "| `bshd` | `bs3hd` | `bsh3d` | `bshd_bs2hd` | `bshd_bsh2d` | `bshd_bshd_bshd` |\n", - "| `thd` | `t3hd` | `th3d` | `thd_t2hd` | `thd_th2d` | `thd_thd_thd` |\n", - "\n", - "The notation system is that `b` stands for the batch size, `s` sequence length, `h` number of attention heads, `d` head dimension, and `t` the total number of tokens in the batch, i.e. `t = sum(s_i) for i in 0,...,b-1`. Here are a few examples of the layouts and their explanations to help clarify the definition.\n", - "\n", - "**qkv_layout=sb3hd:**\n", - "`q`, `k`, `v` are sequence first, i.e. `s` is the leading dimension in each tensor. They are different slices of one tensor `qkv`: `q, k, v = [qkv[:,:,i,:,:] for i in range(3)]`. They are interleaved at the `h * d` dimension.\n", - "\n", - "**qkv_layout=bshd_bsh2d:**\n", - "`q`, `k`, `v` are batch first, i.e. `b` is the leading dimension in each tensor. `q` is contiguous, and `k`, `v` are different slices of tensor `kv`: `k, v = [kv[:,:,:,i,:] for i in range(2)]`. `k`, `v` are interleaved at the `d` dimension.\n", - "\n", - "The `s` and `h` in `bsh2d` are the max sequence length and number of heads for `k`, `v`, which can be different from the `s` and `h` in `bshd` for `q`. We denoted them as the same for brevity reasons. Transformer Engine does differentiate their values for actual execution.\n", - "\n", - "**qkv_layout=thd_thd_thd:**\n", - "`q`, `k`, `v` have variable sequence lengths in a batch. They are all contiguous and have no interleaving.\n", - "\n", - "As of v2.0, Transformer Engine has the following support matrix.\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    BackendSupported QKV FormatsNotes
    flash-attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
    cuDNN attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
    \n", - " JAX: `bs3hd`, `bshd_bs2hd`, `bshd_bshd_bshd` layouts\n", - "
    Framework-native attention`bshd`, `sbhd`PyTorch, JAX: 2 formats, i.e. 10 layouts
    \n", - "\n", - "Some example usage of the different layouts can be found at [test_dpa_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_dpa_qkv_layout_thd](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). Transformer Engine also provides a utility function [transformer_engine.pytorch.attention.dot_product_attention.utils.get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) to help determine which layout a set of `q`, `k`, `v` tensors have (PyTorch only).\n", - "\n", - "
    \n", - "Note\n", - " \n", - "When RoPE is employed, the qkv_layout may change in Transformer Engine PyTorch through [get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py). This is due to the in-place nature of our RoPE implementations. We convert `q`, `k`, `v` tensors from their initial layout to the corresponding hd_hd_hd layout. For example, from sbh3d in pytorch.MultiHeadAttention before RoPE, to sbhd_sbhd_sbhd in pytorch.DotProductAttention after RoPE.\n", - "
    \n" - ], - "id": "fbdcb327" - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 3.2 Attention Mask\n", - "\n", - "Transformer Engine supports 7 mask types, and all the masks are defined as `True` masking out the corresponding element and `False` including the corresponding element in attention calculation.\n", - "\n", - "- `no_mask`, `padding`, `causal`, `causal_bottom_right`, `padding_causal`, `padding_causal_bottom_right`, `arbitrary`\n", - "\n", - "Different backends offer different support for attention mask. As of Transformer Engine 2.0,\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    BackendSupported Mask TypesRequires `attention_mask`
    flash-attention
  • `no_mask`, `causal` (self-attention),
  • `padding`, `padding_causal` (self-attention),
  • `causal_bottom_right`, `padding_causal_bottom_right`
  • `no_mask`, `causal` `causal_bottom_right`: No
  • `padding`, `padding_causal`, `padding_causal_bottom_right`: Yes if `cu_seqlens` not provided
  • `arbitrary`: Yes
    • cuDNN attention
  • `no_mask`, `causal`,
  • `padding`, `padding_causal`,
  • `causal_bottom_right`, `padding_causal_bottom_right`
    • Framework-native attention
  • All (PyTorch)
  • `no_mask`, `causal`, `padding` (Jax)
    • \n", - "\n", - "**Padding masks:** For `padding`, `padding_causal`, `padding_causal_bottom_right` mask types, users need to provide sequence length information to help Transformer Engine figure out where each sequence ends in a batch. As of Transformer Engine 2.0, there are two options to do so in PyTorch and one in JAX.\n", - "\n", - "* PyTorch: When both options are provided by the user, `cu_seqlens` is preferred as there is no extra conversion needed.\n", - " - `cu_seqlens`: Users can provide cumulative sequence length tensors `cu_seqlens_q` and `cu_seqlens_kv` for `q` and `k`/`v` to the flash-attention or cuDNN attention backend. An example of `cu_seqlens` is `[0, 2, 6, 7]` for a batch of 3 `[aa000, bbbb0, c0000]`.\n", - " - `attention_mask`: Users can also provide `attention_mask` as an alternative, which will then be converted to `cu_seqlens`. For self-attention, `attention_mask` should be one single tensor of shape `[batch_size, 1, 1, seqlen_q]`, and for cross-attention, `attention_mask` should be a list of two tensors of shapes `[batch_size, 1, 1, seqlen_q]` and `[batch_size, 1, 1, seqlen_kv]`, respectively.\n", - "\n", - "\n", - "* JAX: Users should provide the `attention_mask` tensor of shape `[batch_size, 1, seqlen_q, seqlen_kv]`.\n", - "\n", - "**qkv_format=thd:** Transformer Engine extracts the max sequence length information from `q`, `k`, `v` if `max_seqlen_q` and `max_seqlen_kv` are not provided. This requires GPU-CPU copy and synchronization operations. For performance reasons, please set `max_seqlen_q` and `max_seqlen_kv` to their appropriate values for `thd` QKV format.\n", - "\n", - "**Arbitrary mask:** cuDNN does not support `Arbitrary` mask type as of v9.3. However, users can convert the mask to a regular `post_scale_bias` bias and achieve the same functionality. An example script for this conversion is [arbitrary_mask_to_post_scale_bias.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py).\n" - ], - "id": "855d9616" - }, + "cells": [ + { + "cell_type": "markdown", + "id": "040f466a", + "metadata": {}, + "source": [ + "# Attention Is All You Need!\n", + "\n", + "The core idea behind Transformer models is the attention mechanism [[1]](https://arxiv.org/abs/1706.03762). It identifies the correlation between words, selects the most important parts of the sentence to focus on, and captures meaningful patterns and dependencies in the data. Figure 1 shows a typical attention mechanism, where pre-softmax operations can be a combination of scaling, bias and masking while the post-softmax operation is often just dropout.\n", + "\n", + "
    \n", + "\n", + "
    Figure 1: Dot product attention.
    \n", + "
    \n", + "\n", + "[Transformer Engine](https://github.com/NVIDIA/TransformerEngine.git) supports the calculation of dot product attention in two frameworks, [PyTorch](https://github.com/pytorch/pytorch) and [JAX](https://github.com/google/jax). The API for each framework is\n", + "\n", + "- [transformer_engine.pytorch.DotProductAttention](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention)\n", + "- [transformer_engine.jax.flax.DotProductAttention](../../api/jax.rst#transformer_engine.jax.flax.DotProductAttention)" + ] + }, + { + "cell_type": "markdown", + "id": "89a7d849", + "metadata": {}, + "source": [ + "## 1. Attention Backends\n", + "\n", + "Transformer Engine provides multiple attention backends for each supported framework. The framework-native backends provide a robust baseline, while the fused, GPU-optimized implementations offer more performance. For example, the flash-attention and cuDNN attention backends in PyTorch. The framework-native backends are often named with \"unfused\", while the more optimized backends are \"fused\" or \"flash\".\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    FrameworkBackend (Module Name)Module Location
    PyTorchcuDNN attention (`FusedAttention`) [transformer_engine.pytorch.attention](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py)
    flash-attention (`FlashAttention`)
    \n", + " PyTorch-native attention (`UnfusedDotProductAttention`)\n", + "
    JAXcuDNN attention (`_FusedDotProductAttention`)[transformer_engine.jax.flax.transformer](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/flax/transformer.py)
    JAX-native attention (`_UnfusedDotProductAttention`)
    " + ] + }, + { + "cell_type": "markdown", + "id": "c90a2573", + "metadata": {}, + "source": [ + "### 1.1 Flash vs. Non-Flash\n", + "\n", + "The attention calculation has quadratic computational and memory complexities to the sequence length. Its runtime and memory requirements quadruple, when the sequence length doubles. This presents a significant challenge to scale Transformer models up for longer contexts, in order to achieve higher model quality.\n", + "\n", + "Compared to the standard, non-flash algorithm, the flash algorithm [[2]](https://arxiv.org/abs/2205.14135) was proposed to reduce the memory scaling to linear and improve the computational efficiency through optimized memory accesses. It employs the following two distinctive techniques.\n", + "\n", + "- **Tiling:** The non-flash algorithm tries to process the query, key, value tensors in one single step, requiring large amounts of global memory and incurring high volumes of reads/writes between global memory and shared memory. The flash algorithm decomposes the input into several tiles, based on the available shared memory and register size, and it computes the softmax one tile at a time.\n", + "\n", + "- **Recomputation:** The non-flash algorithm stores the softmax matrix (quadratic to sequence length) to global memory for the backward pass, while the flash algorithm only saves the softmax normalization factors (linear to sequence length). This reduces the amount of memory required as well as the bandwidth utilization between global memory and shared memory. Even though there is extra computation incurred in order to recalculate the attention in the backward pass, the bandwidth savings still provide significant improvement in efficiency.\n", + "\n", + "
    \n", + "Note: \n", + " \n", + "Transformer Engine's flash-attention backend, available in PyTorch, and cuDNN attention backend (sub-backends 1 and 2), available in PyTorch and JAX, are both based on the flash algorithm.\n", + "
    \n" + ] + }, + { + "cell_type": "markdown", + "id": "b5ce567d", + "metadata": {}, + "source": [ + "### 1.2 flash-attention\n", + "\n", + "The flash-attention backend, available only in PyTorch, is a module wrapped around the public `flash-attn` package [[3]](https://github.com/Dao-AILab/flash-attention). \n", + "\n", + "The flash-attention backend supports `flash-attn`'s features as well as a few extra functionalities to facilitate the use of `flash-attn`, such as converting the `attention_mask` to cumulative sequence lengths `cu_seqlens` for `padding` mask use cases. Please see `transformer_engine.pytorch.attention.FlashAttention` for details.\n", + "\n", + "The `flash-attn` dependency is regularly updated in Transformer Engine. As of v2.0, Transformer Engine supports `flash-attn` 2.0.6+ (see [setup.py](https://github.com/NVIDIA/TransformerEngine/blob/main/setup.py)).\n", + "\n", + "To understand `flash-attn`'s performance, please refer to their benchmarks [here](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#performance).\n", + "\n", + "### 1.3 cuDNN Attention\n", + "\n", + "The cuDNN attention backend, available in PyTorch and JAX, offers another high-performance solution to the attention calculation. It requires [cuDNN](https://developer.nvidia.com/cudnn) to run, and has several sub-backends to support the different precisions and sequence lengths.\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    Sub-BackendAlgorithmPrecisionSequence LengthArchitectureAdditional info
    1FlashBF16/FP16 Any sm80+ [cuDNN](https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#fused-flash-attention-fprop),\n", + " [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention)\n", + "
    2FlashFP8 cuDNN pre-9.0: ≤512 cuDNN pre-9.0: sm90
    cuDNN 9.0+: Any cuDNN 9.0+: sm90+ cuDNN 9.0+: [cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md#scaled-dot-product-attention-fp8)\n", + "
    \n", + "\n", + "The cuDNN attention backend and flash-attention backend have several notable differences. As of Transformer Engine 2.0, cuDNN 9.3 and `flash-attn` 2.4.2,\n", + "\n", + "- flash-attention only supports the PyTorch framework while cuDNN attention supports PyTorch and JAX.\n", + "- flash-attention supports BF16, FP16 precisions while cuDNN attention also supports FP8 (through its sub-backend 2).\n", + "- flash-attention supports `bshd`, `thd` input formats, without any transposes, and `sbhd` format, with transposes, while cuDNN attention supports all three formats without transposes (see Section 3.1 for more details).\n", + "- flash-attention does not support `post_scale_bias`, and cuDNN attention does.\n", + "- flash-attention supports KV-caching and paged attention, and cuDNN attention does not.\n", + "- flash-attention uses bottom right diagonal for `causal` mask in cross attention (see [change log](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#21-change-behavior-of-causal-flag)), and cuDNN attention supports both top left and bottom right.\n", + "- **Sliding window attention (SWA):** flash-attention has SWA(left, right) support for all mask types except top-left causal masks, with or without dropout, and without bias. cuDNN attention supports SWA(left, 0) starting from 9.2 and SWA(left, right) starting from 9.6, without dropout, and with `bias_type=\"no_bias\"`.\n", + "- flash-attention outperforms cuDNN attention on Ampere architectures, and cuDNN attention has 20-50% advantages on Hopper architectures, based on our benchmarks for a number of commonly-used model configurations.\n", + "\n", + "To compare cuDNN attention and flash-attention, users can modify the `model_configs` dictionary in [benchmarks/attention/benchmark_attention.py](https://github.com/NVIDIA/TransformerEngine/blob/main/benchmarks/attention/benchmark_attention.py) to collect performance numbers. The script runs each entry in `model_configs` for `num_iters` times, each time with one forward pass and one backward pass. Both backends are tried, and if one backend does not have support for the specific user input, the runtimes and speedups in the final table would be 0." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c5b8e3d7", + "metadata": {}, + "outputs": [], + "source": [ + "model_configs = {\n", + " # test: b, h, hg, d, sq, skv, p, mask, bias\n", + " \"test_0\": ModelConfig(2, 16, 16, 64, 512, 512, 0.0, \"no_mask\", \"no_bias\"), # short seq\n", + " \"test_1\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"no_bias\"), # longer seq, mask\n", + " \"test_2\": ModelConfig(2, 16, 16, 128, 2048, 2048, 0.0, \"causal\", \"post_scale_bias\"), # bias\n", + " \"test_3\": ModelConfig(2, 32, 4, 128, 8192, 8192, 0.0, \"causal\", \"no_bias\"), # GQA\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50852cb5", + "metadata": {}, + "outputs": [ { - "cell_type": "code", - "metadata": {}, - "source": [ - "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python arbitrary_mask_to_post_scale_bias.py" - ], - "execution_count": null, - "outputs": [ - { - "output_type": "stream", - "name": "stdout", - "text": [ - "Run with post_scale_bias:\n", - "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", - "\n", - "Run with arbitrary mask:\n", - "[INFO | DotProductAttention]: Running with UnfusedDotProductAttention backend\n", - "\n", - "Test passed!\n" - ] - } - ], - "id": "a1f25a9b" - }, + "name": "stdout", + "output_type": "stream", + "text": [ + "Device 0: NVIDIA H100 80GB HBM3 GPU, sm90 compute capability, 79.1GB memory\n", + "Running test_0 with cuDNN attention and flash-attention...\n", + "Running test_1 with cuDNN attention and flash-attention...\n", + "Running test_2 with cuDNN attention...\n", + "Running test_3 with cuDNN attention and flash-attention...\n", + "\n", + " cuDNN fwd+bwd (ms) flash-attn fwd+bwd (ms) cuDNN vs flash speedup\n", + "test_0 0.0340 0.0468 1.3786\n", + "test_1 0.3664 0.5850 1.5968\n", + "test_2 0.9332 0.0000 0.0000\n", + "test_3 7.4875 11.8879 1.5877\n" + ] + } + ], + "source": [ + "!cd ../../../benchmarks/attention/ && python benchmark_attention.py" + ] + }, + { + "cell_type": "markdown", + "id": "9a615119", + "metadata": {}, + "source": [ + "## 2. Backend Selection\n", + "\n", + "Given the various attention backends, Transformer Engine first determines which backends are eligible for the provided inputs and runtime environment, then applies a preference order among the eligible backends. Eligibility is affected by user environment variables, GPU architecture, installed `flash-attn` and cuDNN versions, data type and FP8 recipe, QKV layout, training or inference mode, dropout, and other attention features.\n", + "\n", + "In PyTorch, the candidates are FlashAttention (`flash-attn` v2, v3, or v4), FusedAttention (cuDNN sub-backends), and UnfusedDotProductAttention. Users can disable whole backend families with `NVTE_FLASH_ATTN`, `NVTE_FUSED_ATTN`, or `NVTE_UNFUSED_ATTN`. In JAX, Transformer Engine checks whether a cuDNN fused-attention kernel is available when `NVTE_FUSED_ATTN=1`; otherwise it falls back to the JAX-native implementation.\n", + "\n", + "At a high level, the architecture-specific PyTorch selection order is:\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    FrameworkSelection Order
    PyTorchsm8x (Ampere/Ada): flash-attention > cuDNN attention > PyTorch-native attention
    sm90 (Hopper): cuDNN attention > flash-attention > PyTorch-native attention
    sm100/sm120 (Blackwell): cuDNN attention > flash-attention > PyTorch-native attention
    cuDNN attention: BF16/FP16 uses sub-backend 1 when eligible; FP8 uses sub-backend 2 when enabled and eligible
    JAXcuDNN attention > JAX-native attention
    \n", + "\n", + "Within FlashAttention, TE uses the installed implementation that is supported for the architecture and input. FlashAttention 3 is Hopper-only (`sm90`). FlashAttention 4 supports `sm80`, `sm90`, `sm100`, and `sm120`; on Hopper, TE prefers FlashAttention 3 over FlashAttention 4 when both are installed and eligible. On Blackwell, FlashAttention 4 is the Blackwell-specific flash-attention path when installed and eligible, while FlashAttention 2 can still be eligible depending on the installed version and input configuration.\n", + "\n", + "Within cuDNN FusedAttention, TE asks the fused-attention helper which sub-backend is eligible. Sub-backend 1 is the BF16/FP16 flash-based path when available; sub-backend 2 is the FP8 path when FP8 DPA is enabled and the architecture, cuDNN version, and input configuration support it. Hopper supports eligible FP8 DPA through cuDNN sub-backend 2. In the current PyTorch selector, eligible FP8 DPA on Blackwell is an `sm100` path and is disabled on `sm120`.\n", + "\n", + "When all optimized backends are disabled or ineligible, TE falls back to UnfusedDotProductAttention if it is enabled. If no backend is eligible, backend selection returns no backend and the caller raises an error. As we monitor the performance of different backends, the selection logic may change." + ] + }, + { + "cell_type": "markdown", + "id": "e6c0f3f0", + "metadata": {}, + "source": [ + "### 2.1 Debug Information\n", + "\n", + "To find out which backend is being used during runtime, we have the following two debugging flags. Logging is done by using the `logging` package.\n", + "```\n", + "NVTE_DEBUG = 0/1 # disables/enables debugging\n", + "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n", + "```\n", + "
    \n", + "Note:\n", + " \n", + "These flags are supported in PyTorch only as of Transformer Engine 2.0. JAX support is expected to be added in the future.\n", + "
    " + ] + }, + { + "cell_type": "markdown", + "id": "16660323", + "metadata": {}, + "source": [ + "The example script [example_attention.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/example_attention.py) runs a very basic model with two attention backends, cuDNN attention and flash-attention. Here `NVTE_DEBUG_LEVEL=1` allows us to find out which backend/sub-backend is used in runtime." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "906b8cf1", + "metadata": {}, + "outputs": [ { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Some more examples of running Transformer Engine with different attention masks can be found at [test_dpa_mask](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py).\n", - "\n", - "### 3.3 Attention Bias\n", - "\n", - "Transformer Engine supports 4 attention bias types, `no_bias`, `pre_scale_bias`, `post_scale_bias`, and `ALiBi` (with/without custom slopes). As of Transformer Engine 2.0, their support matrix is as follows.\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
    BackendBias TypeBias ShapeBias Data TypeArchitecture
    flash-attention`no_bias`, `ALiBi` (with slopes)N/AALiBi slopes: FP32sm80+
    cuDNN attentionPyTorch: `no_bias`, `post_scale_bias`, `ALiBi` (without slopes)`post_scale_bias`: BHSS, 1HSS, B1SS, 11SS for forward, 1HSS for backward`post_scale_bias`: same as QKV typecuDNN 8.9.6+: sm90
    JAX: `no_bias`, `post_scale_bias`ALiBi slopes: FP32cuDNN 9.0+: sm80+
    Framework-native attention`no_bias`, `pre_scale_bias`, `post_scale_bias``post_scale_bias`: BHSS, 1HSS, B1SS, 11SS `post_scale_bias`: same as QKV typesm80+
    \n", - "\n", - "The flash-attention backend enables `ALiBi` by asking user to pass in an `alibi_slopes` tensor, which can be the default slopes of vanilla ALiBi, or user-defined slopes. On the other hand, cuDNN attention supports `ALiBi` by taking in a `Boolean` flag, and it only supports vanilla ALiBi as of cuDNN 9.0.\n", - "\n", - "The framework-native backends do not explicitly support `ALiBi`, but users can convert `ALiBi` to a regular `post_scale_bias` bias to achieve the same effect. In PyTorch, this utility function, `transformer_engine.pytorch.attention.get_alibi`, can be used to help with the conversion.\n", - "\n", - "More examples of how to use the various attention biases are at [test_dpa_bias](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)." - ], - "id": "dda4a589" - }, + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Run cuDNN attention...\n", + "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", + "\n", + "Run flash-attention...\n", + "[INFO | DotProductAttention]: Running with FlashAttention backend\n", + "\n", + "Test passed.\n" + ] + } + ], + "source": [ + "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python example_attention.py" + ] + }, + { + "cell_type": "markdown", + "id": "8ca99461", + "metadata": {}, + "source": [ + "`NVTE_DEBUG_LEVEL=2` allows us to find out more about the backend selection logic. Users are encouraged to double check the `config` and provide it to the Transformer Engine team if they would like to file a bug. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3637094", + "metadata": {}, + "outputs": [ { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 3.4 FP8 Attention\n", - "\n", - "A unique feature of Transformer Engine is its FP8 support, not only for the `Linear` layers but also for dot product attention. Transformer Engine's FP8 attention support is through its cuDNN attention sub-backend 2. Recall Figure 1: the two `MatMul` operations are performed in FP8 for computational efficiency, and the `SoftMax` operation is performed in FP32 for numerical accuracy.\n", - "\n", - "Transformer Engine supports FP8 attention through its [C APIs](../../api/c/fused_attn.rst), and [PyTorch API](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention), as of v2.0. Its PyTorch API offers two options, both controlled through the FP8 recipe definition, `transformer_engine.common.recipe.DelayedScaling`.\n", - "\n", - "- `DelayedScaling.fp8_dpa=True (default=False)`: This enables the use of cuDNN attention sub-backend 2, when it does support the provided user inputs. The `FusedAttention` module for cuDNN attention takes FP16 or BF16 tensors as inputs, performs dot product attention in FP8, and returns attention logits in FP16 or BF16 (same as the input type). Casting operations are required to cast tensors to FP8 at the beginning, and back to FP16/BF16 at the end of the module.\n", - "\n", - "- `DelayedScaling.fp8_mha=True (default=False)`: This option, on top of `fp8_dpa=True`, removes the casting operations at the beginning and end of the `FusedAttention` module. This feature is experimental. \n", - "\n", - "Examples of using the two features are available at [test_dpa_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_mha_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). To disable FP8 attention for backward and only use it for forward, users can also set `NVTE_FP8_DPA_BWD=0 (default=1)`." - ], - "id": "a0702339" + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Run cuDNN attention...\n", + "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", + "[DEBUG | DotProductAttention]: Disabling FlashAttention due to NVTE_FLASH_ATTN=0\n", + "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=False, FusedAttention=True (sub-backend 1), UnfusedDotProductAttention=True}\n", + "[DEBUG | DotProductAttention]: Selected backend = FusedAttention (sub-backend 1)\n", + "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", + "\n", + "Run flash-attention...\n", + "[DEBUG | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': , 'cudnn_version': '9.3.0', 'qkv_type': , 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}\n", + "[DEBUG | DotProductAttention]: Disabling FusedAttention due to NVTE_FUSED_ATTN=0\n", + "[DEBUG | DotProductAttention]: Available backends = {FlashAttention=True, FusedAttention=False, UnfusedDotProductAttention=True}\n", + "[DEBUG | DotProductAttention]: Selected backend = FlashAttention\n", + "[INFO | DotProductAttention]: Running with FlashAttention backend\n", + "\n", + "Test passed.\n" + ] } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.12" + ], + "source": [ + "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py" + ] + }, + { + "cell_type": "markdown", + "id": "611d8fdb", + "metadata": {}, + "source": [ + "### 2.2 User Control\n", + "\n", + "Users usually do not need to worry about the backend selection. However, if there is a convergence or performance issue encountered, Transformer Engine provides a few other environment variables for users to experiment with different backends.\n", + "\n", + "**flash-attention or cuDNN attention:**\n", + "Users can enable/disable the flash-attention backend or cuDNN attention backend via the following two environment variables in PyTorch.\n", + "```\n", + "NVTE_FLASH_ATTN = 0 # disables flash-attention; default = 1\n", + "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n", + "```\n", + "\n", + "```\n", + "
    \n", + "Note\n", + " \n", + "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", + "
    \n", + "\n", + "### 2.3 Example Tests\n", + "\n", + "Our [unit tests](https://github.com/NVIDIA/TransformerEngine/tree/main/tests) demonstrate the use of Transformer Engine dot product attention APIs. Users are encouraged to use them as a template when integrating Transformer Engine to their ML workflows.\n", + "\n", + "For example, in PyTorch, [test_dot_product_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) offers a variety of use cases of `pytorch.DotProductAttention`, from data types, model configs, checkpointing, to QKV layouts." + ] + }, + { + "cell_type": "markdown", + "id": "e60a2a3e", + "metadata": {}, + "source": [ + "## 3. Backend Support\n", + "\n", + "Transformer Engine supports commonly-used features such as self and cross attention, FP16/BF16 precisions, dropout, and checkpointing. But it also offers a range of other features. As of v2.0, Transformer Engine's attention backends have the following support matrix.\n", + "\n", + "| Attention Backend | Precision | Architecture | Sliding Window Attention | MQA/GQA | Multi-Latent Attention | Context Parallelism | Determinism Possible |\n", + "| :---------------- | :-------- | :----------- | :----------------------- | :------ | :--------------------- | :------------------ | :------------ |\n", + "| cuDNN attention (all frameworks) | BF16, FP16, FP8 (PyTorch only) | sm80+ | Yes (cuDNN 9.2+) | Yes | Yes | Yes (`bshd`,`sbhd`, `thd`) | Yes |\n", + "| flash-attention (PyTorch) | BF16, FP16 | sm80+ | Yes | Yes | Yes | Yes (`bshd`,`thd`) | Yes |\n", + "| Framework-native attention | BF16, FP16, FP32 | Any | No, unless used as a mask | Yes | Yes (PyTorch only) | No | Yes |\n", + "\n", + "Some unit tests are provided to serve as a starting point for integrating such features into users' models. For example,\n", + "- sliding window attention: [test_dpa_swa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", + "- MQA/GQA: [test_te_layer_mqa_gqa](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", + "- Multi-Latent Attention: [test_dpa_mla](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)\n", + "- context parallelism: [test_cp_with_fused_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py), [test_cp_with_flash_attention](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention_with_cp.py)" + ] + }, + { + "cell_type": "markdown", + "id": "fbdcb327", + "metadata": {}, + "source": [ + "### 3.1 QKV Layout\n", + "\n", + "Transformer Engine supports various layouts of the query `q`, key `k`, value `v` tensors. It has defined 15 QKV layouts, which are grouped into 3 QKV formats and 5 QKV layout groups to help with similar memory/computational operations across different layouts. The mapping relationships of these layouts and groups are,\n", + "\n", + "| `qkv_layout`         | `qkv_layout_group`=`3hd` | `h3d` | `hd_2hd` | `hd_h2d` | `hd_hd_hd` |\n", + "| ----------: | -----------: | -----: | ----------: | ----------: | -------------: |\n", + "| `qkv_format`=`sbhd` | `sb3hd` | `sbh3d` | `sbhd_sb2hd` | `sbhd_sbh2d` | `sbhd_sbhd_sbhd` |\n", + "| `bshd` | `bs3hd` | `bsh3d` | `bshd_bs2hd` | `bshd_bsh2d` | `bshd_bshd_bshd` |\n", + "| `thd` | `t3hd` | `th3d` | `thd_t2hd` | `thd_th2d` | `thd_thd_thd` |\n", + "\n", + "The notation system is that `b` stands for the batch size, `s` sequence length, `h` number of attention heads, `d` head dimension, and `t` the total number of tokens in the batch, i.e. `t = sum(s_i) for i in 0,...,b-1`. Here are a few examples of the layouts and their explanations to help clarify the definition.\n", + "\n", + "**qkv_layout=sb3hd:**\n", + "`q`, `k`, `v` are sequence first, i.e. `s` is the leading dimension in each tensor. They are different slices of one tensor `qkv`: `q, k, v = [qkv[:,:,i,:,:] for i in range(3)]`. They are interleaved at the `h * d` dimension.\n", + "\n", + "**qkv_layout=bshd_bsh2d:**\n", + "`q`, `k`, `v` are batch first, i.e. `b` is the leading dimension in each tensor. `q` is contiguous, and `k`, `v` are different slices of tensor `kv`: `k, v = [kv[:,:,:,i,:] for i in range(2)]`. `k`, `v` are interleaved at the `d` dimension.\n", + "\n", + "The `s` and `h` in `bsh2d` are the max sequence length and number of heads for `k`, `v`, which can be different from the `s` and `h` in `bshd` for `q`. We denoted them as the same for brevity reasons. Transformer Engine does differentiate their values for actual execution.\n", + "\n", + "**qkv_layout=thd_thd_thd:**\n", + "`q`, `k`, `v` have variable sequence lengths in a batch. They are all contiguous and have no interleaving.\n", + "\n", + "As of v2.0, Transformer Engine has the following support matrix.\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    BackendSupported QKV FormatsNotes
    flash-attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
    cuDNN attention`bshd`, `sbhd`, `thd`PyTorch: 3 formats, i.e. 15 layouts
    \n", + " JAX: `bs3hd`, `bshd_bs2hd`, `bshd_bshd_bshd` layouts\n", + "
    Framework-native attention`bshd`, `sbhd`PyTorch, JAX: 2 formats, i.e. 10 layouts
    \n", + "\n", + "Some example usage of the different layouts can be found at [test_dpa_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_dpa_qkv_layout_thd](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). Transformer Engine also provides a utility function [transformer_engine.pytorch.attention.dot_product_attention.utils.get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) to help determine which layout a set of `q`, `k`, `v` tensors have (PyTorch only).\n", + "\n", + "
    \n", + "Note\n", + " \n", + "When RoPE is employed, the qkv_layout may change in Transformer Engine PyTorch through [get_qkv_layout](https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py). This is due to the in-place nature of our RoPE implementations. We convert `q`, `k`, `v` tensors from their initial layout to the corresponding hd_hd_hd layout. For example, from sbh3d in pytorch.MultiHeadAttention before RoPE, to sbhd_sbhd_sbhd in pytorch.DotProductAttention after RoPE.\n", + "
    \n" + ] + }, + { + "cell_type": "markdown", + "id": "855d9616", + "metadata": {}, + "source": [ + "### 3.2 Attention Mask\n", + "\n", + "Transformer Engine supports 7 mask types, and all the masks are defined as `True` masking out the corresponding element and `False` including the corresponding element in attention calculation.\n", + "\n", + "- `no_mask`, `padding`, `causal`, `causal_bottom_right`, `padding_causal`, `padding_causal_bottom_right`, `arbitrary`\n", + "\n", + "Different backends offer different support for attention mask. As of Transformer Engine 2.0,\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    BackendSupported Mask TypesRequires `attention_mask`
    flash-attention
  • `no_mask`, `causal` (self-attention),
  • `padding`, `padding_causal` (self-attention),
  • `causal_bottom_right`, `padding_causal_bottom_right`
  • `no_mask`, `causal` `causal_bottom_right`: No
  • `padding`, `padding_causal`, `padding_causal_bottom_right`: Yes if `cu_seqlens` not provided
  • `arbitrary`: Yes
    • cuDNN attention
  • `no_mask`, `causal`,
  • `padding`, `padding_causal`,
  • `causal_bottom_right`, `padding_causal_bottom_right`
    • Framework-native attention
  • All (PyTorch)
  • `no_mask`, `causal`, `padding` (Jax)
    • \n", + "\n", + "**Padding masks:** For `padding`, `padding_causal`, `padding_causal_bottom_right` mask types, users need to provide sequence length information to help Transformer Engine figure out where each sequence ends in a batch. As of Transformer Engine 2.0, there are two options to do so in PyTorch and one in JAX.\n", + "\n", + "* PyTorch: When both options are provided by the user, `cu_seqlens` is preferred as there is no extra conversion needed.\n", + " - `cu_seqlens`: Users can provide cumulative sequence length tensors `cu_seqlens_q` and `cu_seqlens_kv` for `q` and `k`/`v` to the flash-attention or cuDNN attention backend. An example of `cu_seqlens` is `[0, 2, 6, 7]` for a batch of 3 `[aa000, bbbb0, c0000]`.\n", + " - `attention_mask`: Users can also provide `attention_mask` as an alternative, which will then be converted to `cu_seqlens`. For self-attention, `attention_mask` should be one single tensor of shape `[batch_size, 1, 1, seqlen_q]`, and for cross-attention, `attention_mask` should be a list of two tensors of shapes `[batch_size, 1, 1, seqlen_q]` and `[batch_size, 1, 1, seqlen_kv]`, respectively.\n", + "\n", + "\n", + "* JAX: Users should provide the `attention_mask` tensor of shape `[batch_size, 1, seqlen_q, seqlen_kv]`.\n", + "\n", + "**qkv_format=thd:** Transformer Engine extracts the max sequence length information from `q`, `k`, `v` if `max_seqlen_q` and `max_seqlen_kv` are not provided. This requires GPU-CPU copy and synchronization operations. For performance reasons, please set `max_seqlen_q` and `max_seqlen_kv` to their appropriate values for `thd` QKV format.\n", + "\n", + "**Arbitrary mask:** cuDNN does not support `Arbitrary` mask type as of v9.3. However, users can convert the mask to a regular `post_scale_bias` bias and achieve the same functionality. An example script for this conversion is [arbitrary_mask_to_post_scale_bias.py](https://raw.githubusercontent.com/NVIDIA/TransformerEngine/main/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1f25a9b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Run with post_scale_bias:\n", + "[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)\n", + "\n", + "Run with arbitrary mask:\n", + "[INFO | DotProductAttention]: Running with UnfusedDotProductAttention backend\n", + "\n", + "Test passed!\n" + ] } + ], + "source": [ + "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python arbitrary_mask_to_post_scale_bias.py" + ] + }, + { + "cell_type": "markdown", + "id": "dda4a589", + "metadata": {}, + "source": [ + "Some more examples of running Transformer Engine with different attention masks can be found at [test_dpa_mask](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py).\n", + "\n", + "### 3.3 Attention Bias\n", + "\n", + "Transformer Engine supports 4 attention bias types, `no_bias`, `pre_scale_bias`, `post_scale_bias`, and `ALiBi` (with/without custom slopes). As of Transformer Engine 2.0, their support matrix is as follows.\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
    BackendBias TypeBias ShapeBias Data TypeArchitecture
    flash-attention`no_bias`, `ALiBi` (with slopes)N/AALiBi slopes: FP32sm80+
    cuDNN attentionPyTorch: `no_bias`, `post_scale_bias`, `ALiBi` (without slopes)`post_scale_bias`: BHSS, 1HSS, B1SS, 11SS for forward, 1HSS for backward`post_scale_bias`: same as QKV typecuDNN 8.9.6+: sm90
    JAX: `no_bias`, `post_scale_bias`ALiBi slopes: FP32cuDNN 9.0+: sm80+
    Framework-native attention`no_bias`, `pre_scale_bias`, `post_scale_bias``post_scale_bias`: BHSS, 1HSS, B1SS, 11SS `post_scale_bias`: same as QKV typesm80+
    \n", + "\n", + "The flash-attention backend enables `ALiBi` by asking user to pass in an `alibi_slopes` tensor, which can be the default slopes of vanilla ALiBi, or user-defined slopes. On the other hand, cuDNN attention supports `ALiBi` by taking in a `Boolean` flag, and it only supports vanilla ALiBi as of cuDNN 9.0.\n", + "\n", + "The framework-native backends do not explicitly support `ALiBi`, but users can convert `ALiBi` to a regular `post_scale_bias` bias to achieve the same effect. In PyTorch, this utility function, `transformer_engine.pytorch.attention.get_alibi`, can be used to help with the conversion.\n", + "\n", + "More examples of how to use the various attention biases are at [test_dpa_bias](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py)." + ] + }, + { + "cell_type": "markdown", + "id": "a0702339", + "metadata": {}, + "source": [ + "### 3.4 FP8 Attention\n", + "\n", + "A unique feature of Transformer Engine is its FP8 support, not only for the `Linear` layers but also for dot product attention. Transformer Engine's FP8 attention support is through its cuDNN attention sub-backend 2. Recall Figure 1: the two `MatMul` operations are performed in FP8 for computational efficiency, and the `SoftMax` operation is performed in FP32 for numerical accuracy.\n", + "\n", + "Transformer Engine supports FP8 attention through its [C APIs](../../api/c/fused_attn.rst), and [PyTorch API](../../api/pytorch.rst#transformer_engine.pytorch.DotProductAttention), as of v2.0. Its PyTorch API offers two options, both controlled through the FP8 recipe definition, `transformer_engine.common.recipe.DelayedScaling`.\n", + "\n", + "- `DelayedScaling.fp8_dpa=True (default=False)`: This enables the use of cuDNN attention sub-backend 2, when it does support the provided user inputs. The `FusedAttention` module for cuDNN attention takes FP16 or BF16 tensors as inputs, performs dot product attention in FP8, and returns attention logits in FP16 or BF16 (same as the input type). Casting operations are required to cast tensors to FP8 at the beginning, and back to FP16/BF16 at the end of the module.\n", + "\n", + "- `DelayedScaling.fp8_mha=True (default=False)`: This option, on top of `fp8_dpa=True`, removes the casting operations at the beginning and end of the `FusedAttention` module. This feature is experimental. \n", + "\n", + "Examples of using the two features are available at [test_dpa_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py) and [test_mha_fp8_vs_f16](https://github.com/NVIDIA/TransformerEngine/blob/main/tests/pytorch/attention/test_attention.py). To disable FP8 attention for backward and only use it for forward, users can also set `NVTE_FP8_DPA_BWD=0 (default=1)`." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From cbf6550af5d6789c13cde9d3a51ac97500362714 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 24 Jul 2026 04:00:28 -0700 Subject: [PATCH 54/63] add fused attn graph cache debug code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 6 + .../fused_attn_f16_arbitrary_seqlen.cu | 41 ++- .../common/fused_attn/fused_attn_fp8.cu | 41 ++- .../common/fused_attn/graph_cache_debug.h | 236 ++++++++++++++++++ 4 files changed, 304 insertions(+), 20 deletions(-) create mode 100644 transformer_engine/common/fused_attn/graph_cache_debug.h diff --git a/docs/envvars.rst b/docs/envvars.rst index bf32df8971..d6ffa241ef 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -177,6 +177,12 @@ backend-selection overview. :Default: ``0`` :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. +.. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). When set to ``1``, prints to stderr (prefixed ``[FUSED-ATTN-CACHE]``) a per-lookup ``HIT``/``MISS`` line with the full graph-cache key, a ``BUILD`` line whenever a new graph is constructed, an ``EXEC`` line whenever a graph is executed, a ``SUMMARY`` of graph builds vs. executions at process exit, and a breakdown of cuDNN graph-build timings. Useful for diagnosing redundant graph rebuilds or stale-cache reuse, and for profiling graph-build cost. Has negligible overhead when unset. + .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO :Type: ``int`` (0 or 1) 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 84c46dcdd1..13806aa5dc 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 @@ -18,6 +18,7 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" +#include "graph_cache_debug.h" #include "utils.h" namespace transformer_engine { @@ -167,6 +168,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } + graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); if (cache_hit) { return cached_graph; } @@ -434,16 +436,24 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer( + "fwd", graph_cache_debug::BuildStage::BuildOpGraph, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); + graph_cache_debug::timer( + "fwd", graph_cache_debug::BuildStage::CreatePlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); + graph_cache_debug::record_build("fwd"); // Lock the insert. If another thread inserted a graph for the same key while we were building, // use their graph (it's the same as ours) and discard our graph. { @@ -484,6 +494,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } + graph_cache_debug::record_exec("fwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -712,6 +723,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } + graph_cache_debug::record_cache_lookup("bwd", cache_hit, cfg); if (cache_hit) { return cached_graph; } @@ -951,15 +963,23 @@ void fused_attn_arbitrary_seqlen_bwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer( + "bwd", graph_cache_debug::BuildStage::BuildOpGraph, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); + graph_cache_debug::timer( + "bwd", graph_cache_debug::BuildStage::CreatePlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); + graph_cache_debug::record_build("bwd"); // Lock the insert. If another thread inserted a graph for the same key while we were building, // use their graph (it's the same as ours) and discard our graph. { @@ -995,6 +1015,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } + graph_cache_debug::record_exec("bwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 37082ed9cf..531034aecc 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -11,6 +11,7 @@ #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" +#include "graph_cache_debug.h" #include "utils.h" namespace transformer_engine { @@ -142,6 +143,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } + graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); if (cache_hit) { return cached_graph; } @@ -393,14 +395,22 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer( + "fwd", graph_cache_debug::BuildStage::BuildOpGraph, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); + graph_cache_debug::timer( + "fwd", graph_cache_debug::BuildStage::CreatePlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); + graph_cache_debug::record_build("fwd"); // Lock the insert. If another thread inserted a graph for the same key while we were building, // use their graph (it's the same as ours) and discard our graph. { @@ -424,6 +434,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } + graph_cache_debug::record_exec("fwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -628,6 +639,7 @@ void fused_attn_fp8_bwd_impl( cache_hit = (it != cache.end()); if (cache_hit) cached_graph = it->second; } + graph_cache_debug::record_cache_lookup("bwd", cache_hit, cfg); if (cache_hit) { return cached_graph; } @@ -1008,15 +1020,23 @@ void fused_attn_fp8_bwd_impl( auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer( + "bwd", graph_cache_debug::BuildStage::BuildOpGraph, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); + graph_cache_debug::timer( + "bwd", graph_cache_debug::BuildStage::CreatePlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); + graph_cache_debug::record_build("bwd"); // Lock the insert. If another thread inserted a graph for the same key while we were building, // use their graph (it's the same as ours) and discard our graph. { @@ -1039,6 +1059,7 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } + graph_cache_debug::record_exec("bwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h new file mode 100644 index 0000000000..8df0e6654f --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -0,0 +1,236 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// ============================================================================ +// Fused-attention graph cache diagnostics. +// +// Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG=1 to get the cache event +// counters and graph build timings, to help diagnose redundant graph rebuilds +// or stale-cache reuse, and to profile graph-build cost. +// ============================================================================ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "config_and_params.h" + +namespace transformer_engine { +namespace fused_attn { +namespace graph_cache_debug { + +// Enable diagnostics with NVTE_FUSED_ATTN_CACHE_DEBUG=1. Single read at startup, cached. +// Negligible overhead when unset. +inline bool enabled() { + static const bool on = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return on; +} + +// More readable, shorter thread IDs (0, 1, 2, ...). +inline unsigned thread_seq_id() { + static std::atomic next{0}; + static thread_local unsigned id = next.fetch_add(1, std::memory_order_relaxed); + return id; +} + +// Registered at first use. On process exit, prints overall event counters and +// graph build timings. +inline void register_summary_once(); + +// ============================================================================ +// Cache event counters (forward/backward): +// - BUILD: a successful graph build; triggered by a cache miss +// - EXEC: a graph execution call with valid runtime tensors +// - HIT: a cache lookup that hit; may not trigger an EXEC, and may only be +// a backend availability check or from the first shape-probing call of +// nvte_fused_attn_fwd/bwd which has no runtime tensors +// - MISS: a cache lookup that missed; triggers a graph build +// ============================================================================ + +struct EventCounters { + std::atomic built{0}; + std::atomic exec{0}; + std::atomic hit{0}; + std::atomic miss{0}; +}; + +inline EventCounters &counters(bool is_fwd) { + static EventCounters fwd; + static EventCounters bwd; + return is_fwd ? fwd : bwd; +} + +inline void print_counters(const char *event) { + const EventCounters &f = counters(/*is_fwd=*/true); + const EventCounters &b = counters(/*is_fwd=*/false); + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%llu exec=%llu hit=%llu miss=%llu | " + "bwd built=%llu exec=%llu hit=%llu miss=%llu\n", + event, thread_seq_id(), + static_cast(f.built.load(std::memory_order_relaxed)), + static_cast(f.exec.load(std::memory_order_relaxed)), + static_cast(f.hit.load(std::memory_order_relaxed)), + static_cast(f.miss.load(std::memory_order_relaxed)), + static_cast(b.built.load(std::memory_order_relaxed)), + static_cast(b.exec.load(std::memory_order_relaxed)), + static_cast(b.hit.load(std::memory_order_relaxed)), + static_cast(b.miss.load(std::memory_order_relaxed))); + std::fflush(stderr); +} + +inline void record_build(const char *pass) { + if (!enabled()) return; + register_summary_once(); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + counters(is_fwd).built.fetch_add(1, std::memory_order_relaxed); + print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD"); +} + +inline void record_exec(const char *pass) { + if (!enabled()) return; + register_summary_once(); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); + print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC"); +} + +inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { + if (!enabled()) return; + register_summary_once(); + EventCounters &pc = counters(std::strcmp(pass, "fwd") == 0); + (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] %-3s %-4s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld " + "bias=%lld wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " + "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " + "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " + "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " + "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", + pass, hit ? "HIT" : "MISS", thread_seq_id(), static_cast(c.is_training), + static_cast(c.deterministic), static_cast(c.cuda_graph), + static_cast(c.return_max_logit), static_cast(c.is_forward), + static_cast(c.attn_mask_type), static_cast(c.bias_type), + static_cast(c.window_size_left), static_cast(c.window_size_right), + static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), + static_cast(c.scaling_mode), static_cast(c.dropout), + static_cast(c.attn_scale), static_cast(c.qkv_dtype), + static_cast(c.o_dtype), static_cast(c.do_dtype), + static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), + static_cast(c.o_format), static_cast(c.do_format), + static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), + static_cast(c.do_scale_inv_format), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.head_dim_qk), static_cast(c.head_dim_v), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), + static_cast(c.num_pages_v), static_cast(c.page_size_k), + static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), + static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), + static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), + static_cast(c.bias_seqlen_kv)); + std::fflush(stderr); +} + +// ============================================================================ +// Graph build timings for individual cuDNN-frontend calls in forward/backward: +// e.g. `validate`, `build_operation_graph`, `create_execution_plans`, +// `check_support`, `build_plans` +// ============================================================================ + +enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; +inline constexpr const char *kStageNames[] = {"validate", "build_operation_graph", + "create_execution_plans", "check_support", + "build_plans"}; +struct StageTiming { + std::atomic calls{0}; + std::atomic time_ns{0}; +}; +constexpr size_t kStageBuckets = 2 * static_cast(BuildStage::kCount); +inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { + static std::array table{}; + const size_t idx = + (is_fwd ? 0u : 1u) * static_cast(BuildStage::kCount) + static_cast(s); + return table[idx]; +} + +struct ScopedBuildTimer { + BuildStage stage; + bool on; + bool is_fwd; + std::chrono::steady_clock::time_point start; + ScopedBuildTimer(bool is_fwd_, BuildStage s) : stage(s), on(enabled()), is_fwd(is_fwd_) { + if (!on) return; + register_summary_once(); + start = std::chrono::steady_clock::now(); + } + ~ScopedBuildTimer() { + if (!on) return; + const uint64_t elapsed_ns = static_cast( + std::chrono::duration_cast(std::chrono::steady_clock::now() - + start) + .count()); + StageTiming &t = stage_timing(is_fwd, stage); + t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); + t.calls.fetch_add(1, std::memory_order_relaxed); + } +}; + +template +inline void timer(const char *pass, BuildStage stage, Fn &&fn) { + ScopedBuildTimer scoped(std::strcmp(pass, "fwd") == 0, stage); + fn(); +} + +// ============================================================================ +// Summary: on process exit, print cache event counters and graph build timings. +// ============================================================================ +inline void register_summary_once() { + static const bool registered = [] { + std::atexit([] { + if (!enabled()) return; + print_counters("SUMMARY"); + for (int p = 0; p < 2; ++p) { + const bool is_fwd = (p == 0); + const char *pass = is_fwd ? "fwd" : "bwd"; + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + const BuildStage s = static_cast(i); + const StageTiming &t = stage_timing(is_fwd, s); + const uint64_t n = t.calls.load(std::memory_order_relaxed); + if (n == 0) continue; + const double total_ms = + static_cast(t.time_ns.load(std::memory_order_relaxed)) / 1e6; + std::fprintf(stderr, + "[FUSED-ATTN-CACHE] %-3s %-22s | calls=%llu | time=%9.1f ms | avg=%9.3f ms/call\n", + pass, kStageNames[i], static_cast(n), total_ms, + total_ms / n); + } + } + std::fflush(stderr); + }); + return true; + }(); + (void)registered; +} + +} // namespace graph_cache_debug +} // namespace fused_attn +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ From bf4bfc0d17e9aedc24dfffba31a2e156d3d20060 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 24 Jul 2026 05:34:53 -0700 Subject: [PATCH 55/63] fix kv cache probes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_kv_cache.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index cdd98d2445..ad57d584f3 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -4,6 +4,7 @@ from collections import OrderedDict from typing import List +import copy import os import sys import pathlib @@ -472,8 +473,11 @@ def test_kv_cache(dtype, model, qkv_format, is_paged, backend, module, is_cuda_g qkv_layout = qkv_format + "_" + "_".join([inference_params_qkv_format] * 2) if is_paged: qkv_layout = "paged_kv_" + qkv_layout - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, + # probe inference configs only; reference configs are widely supported + probe_config = copy.deepcopy(config) + probe_config.attn_mask_type = "padding_causal" + available_backends, _, _ = get_available_attention_backends( + probe_config, qkv_dtype=dtype, qkv_layout=qkv_layout, pad_between_seqs=False, From c056cbe1e2f7cc4aeb0d150ff10057af7b064763 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Sun, 26 Jul 2026 12:46:18 -0700 Subject: [PATCH 56/63] deduplicate L0 pytest tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index e63b7b7b04..4752c496f1 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,6 +1,7 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +import copy import logging import os import sys @@ -139,7 +140,7 @@ def test_dot_product_attention( tols = dict(atol=1e-3, rtol=1e-3) if dtype == torch.bfloat16: tols = dict(atol=1.5e-2, rtol=1.5e-2) - config = model_configs[model] + config = copy.deepcopy(model_configs[model]) is_mla = config.head_dim_qk != config.head_dim_v is_mqa_gqa = config.num_heads != config.num_gqa_groups if qkv_layout is None: @@ -550,6 +551,9 @@ def test_dpa_softmax(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_softmax.keys()) def test_dpa_softmax_thd(dtype, model_configs, model): """Test DotProductAttention module with different softmax types""" + config = model_configs[model] + if "padding" not in config.attn_mask_type: + pytest.skip(f"Duplicate test to others with THD and padding mask.") test_dot_product_attention(dtype, model_configs, model, True, "thd_thd_thd", False, False) @@ -820,6 +824,9 @@ def test_dpa_bias_shapes(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "sbhd_sbhd_sbhd"]) def test_dpa_sliding_window(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with sliding window attention""" + config = model_configs[model] + if qkv_layout == "thd_thd_thd" and "padding" not in config.attn_mask_type: + pytest.skip(f"Duplicate test to others with THD and padding mask.") test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, True, False) @@ -1985,6 +1992,9 @@ def test_mha_fp8_vs_f16( scaling_mode, ): """Test MultiHeadAttention module in FP8""" + if not is_training and fp8_dpa_bwd: + pytest.skip("fp8_dpa_bwd=True not applicable for inference") + os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" config = model_configs_fp8_vs_f16[model] @@ -2234,6 +2244,10 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): """Test DotProductAttention module in FP8""" config = model_configs_fp8_vs_f16[model] + if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: + pytest.skip("qkv_layout not applicable for MQA/GQA") + if not is_training and fp8_dpa_bwd: + pytest.skip("fp8_dpa_bwd=True not applicable for inference") # TODO(cyang): think of another way to verify dropout results # test cuDNN FP8 dropout @@ -2293,8 +2307,6 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal pytest.skip("No FP8 attention backend available.") if not fused_attn_supported_f16: pytest.skip("No reference backend available.") - if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: - pytest.skip("qkv_layout not applicable for MQA/GQA") if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" From 16df390afc49c94f96122e61c4d0a934fc67a0f7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 26 Jul 2026 19:52:22 +0000 Subject: [PATCH 57/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../fused_attn_f16_arbitrary_seqlen.cu | 36 +++++++------- .../common/fused_attn/fused_attn_fp8.cu | 36 +++++++------- .../common/fused_attn/graph_cache_debug.h | 49 +++++++++---------- 3 files changed, 60 insertions(+), 61 deletions(-) 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 13806aa5dc..b7c7a349af 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 @@ -437,17 +437,17 @@ void fused_attn_arbitrary_seqlen_fwd_impl( : std::make_tuple(nullptr, nullptr); graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::Validate, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); - graph_cache_debug::timer( - "fwd", graph_cache_debug::BuildStage::BuildOpGraph, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); - graph_cache_debug::timer( - "fwd", graph_cache_debug::BuildStage::CreatePlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CheckSupport, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildPlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, @@ -964,17 +964,17 @@ void fused_attn_arbitrary_seqlen_bwd_impl( : std::make_tuple(nullptr, nullptr); graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::Validate, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); - graph_cache_debug::timer( - "bwd", graph_cache_debug::BuildStage::BuildOpGraph, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); - graph_cache_debug::timer( - "bwd", graph_cache_debug::BuildStage::CreatePlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CheckSupport, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildPlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 531034aecc..9449f74206 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -396,17 +396,17 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de : std::make_tuple(nullptr, nullptr); graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::Validate, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); - graph_cache_debug::timer( - "fwd", graph_cache_debug::BuildStage::BuildOpGraph, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); - graph_cache_debug::timer( - "fwd", graph_cache_debug::BuildStage::CreatePlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CheckSupport, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildPlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); @@ -1021,17 +1021,17 @@ void fused_attn_fp8_bwd_impl( : std::make_tuple(nullptr, nullptr); graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::Validate, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); - graph_cache_debug::timer( - "bwd", graph_cache_debug::BuildStage::BuildOpGraph, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); }); - graph_cache_debug::timer( - "bwd", graph_cache_debug::BuildStage::CreatePlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CheckSupport, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildPlans, - [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 8df0e6654f..843e4156b4 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -76,19 +76,18 @@ inline EventCounters &counters(bool is_fwd) { inline void print_counters(const char *event) { const EventCounters &f = counters(/*is_fwd=*/true); const EventCounters &b = counters(/*is_fwd=*/false); - std::fprintf( - stderr, - "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%llu exec=%llu hit=%llu miss=%llu | " - "bwd built=%llu exec=%llu hit=%llu miss=%llu\n", - event, thread_seq_id(), - static_cast(f.built.load(std::memory_order_relaxed)), - static_cast(f.exec.load(std::memory_order_relaxed)), - static_cast(f.hit.load(std::memory_order_relaxed)), - static_cast(f.miss.load(std::memory_order_relaxed)), - static_cast(b.built.load(std::memory_order_relaxed)), - static_cast(b.exec.load(std::memory_order_relaxed)), - static_cast(b.hit.load(std::memory_order_relaxed)), - static_cast(b.miss.load(std::memory_order_relaxed))); + std::fprintf(stderr, + "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%llu exec=%llu hit=%llu miss=%llu | " + "bwd built=%llu exec=%llu hit=%llu miss=%llu\n", + event, thread_seq_id(), + static_cast(f.built.load(std::memory_order_relaxed)), + static_cast(f.exec.load(std::memory_order_relaxed)), + static_cast(f.hit.load(std::memory_order_relaxed)), + static_cast(f.miss.load(std::memory_order_relaxed)), + static_cast(b.built.load(std::memory_order_relaxed)), + static_cast(b.exec.load(std::memory_order_relaxed)), + static_cast(b.hit.load(std::memory_order_relaxed)), + static_cast(b.miss.load(std::memory_order_relaxed))); std::fflush(stderr); } @@ -138,7 +137,8 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi static_cast(c.head_dim_qk), static_cast(c.head_dim_v), static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), - static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_batch_size), + static_cast(c.bucketed_num_tokens_q), static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), static_cast(c.num_pages_v), static_cast(c.page_size_k), static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), @@ -155,9 +155,8 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi // ============================================================================ enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; -inline constexpr const char *kStageNames[] = {"validate", "build_operation_graph", - "create_execution_plans", "check_support", - "build_plans"}; +inline constexpr const char *kStageNames[] = { + "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; struct StageTiming { std::atomic calls{0}; std::atomic time_ns{0}; @@ -182,10 +181,10 @@ struct ScopedBuildTimer { } ~ScopedBuildTimer() { if (!on) return; - const uint64_t elapsed_ns = static_cast( - std::chrono::duration_cast(std::chrono::steady_clock::now() - - start) - .count()); + const uint64_t elapsed_ns = + static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count()); StageTiming &t = stage_timing(is_fwd, stage); t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); t.calls.fetch_add(1, std::memory_order_relaxed); @@ -216,10 +215,10 @@ inline void register_summary_once() { if (n == 0) continue; const double total_ms = static_cast(t.time_ns.load(std::memory_order_relaxed)) / 1e6; - std::fprintf(stderr, - "[FUSED-ATTN-CACHE] %-3s %-22s | calls=%llu | time=%9.1f ms | avg=%9.3f ms/call\n", - pass, kStageNames[i], static_cast(n), total_ms, - total_ms / n); + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] %-3s %-22s | calls=%llu | time=%9.1f ms | avg=%9.3f ms/call\n", + pass, kStageNames[i], static_cast(n), total_ms, total_ms / n); } } std::fflush(stderr); From a6da26e5619200ace6ed51696957adb645463854 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Sun, 26 Jul 2026 13:22:59 -0700 Subject: [PATCH 58/63] fix merge with torch.compile PRs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../attention/dot_product_attention/utils.py | 37 ++++--------------- 1 file changed, 7 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index b453239b17..79377ea0c0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -428,33 +428,12 @@ def error(self, *args, **kwargs): @torch.compiler.assume_constant_result -def _get_fused_attn_backend( - is_training, - q_type, - kv_type, - qkv_layout, - bias_type, - attn_mask_type, - softmax_type, - *args, -): +def _get_fused_attn_backend(fused_attn_params): """Constant-foldable tex.get_fused_attn_backend: the result depends only on the attention config, and the python-side enum keeps it traceable by - torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax - are taken as their string keys and resolved to the pybind enums here, so - that every argument is a python literal or a python enum.""" - return FusedAttnBackend.cast( - tex.get_fused_attn_backend( - is_training, - q_type, - kv_type, - QKVLayout[qkv_layout], - AttnBiasType[bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], - *args, - ) - ) + torch.compile (see the FusedAttnBackend docstring).""" + fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) + return FusedAttnBackend.cast(fused_attention_backend), reject_message def get_attention_backend( @@ -1640,12 +1619,10 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fused_attn_kwargs["bias_seqlen_kv"] = step_seqlen_kv fused_attn_params = FusedAttentionParams(**fused_attn_kwargs) - # NOTE: under torch.compile the numeric args below must not be symbolic - # (assume_constant_result requires concrete values); ints/floats made + # NOTE: under torch.compile the numeric fields of fused_attn_params must not be + # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. - fused_attention_backend = _get_fused_attn_backend(fused_attn_params) - - fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) + fused_attention_backend, reject_message = _get_fused_attn_backend(fused_attn_params) if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug( "Disabling FusedAttention: %s%s", From e2d1fc971f037815f844e326dc5832a2604720a5 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 27 Jul 2026 08:07:12 -0700 Subject: [PATCH 59/63] remove redundant change Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a5f2b19b21..6c876ad7cc 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2687,7 +2687,6 @@ 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 From fa6e636019721d99c50d43515d672c4ef3dab9db Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 27 Jul 2026 08:08:18 -0700 Subject: [PATCH 60/63] group newly enabled SWA tests to tiers L0/L1 in Jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/jax/test_fused_attn.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index b53ce95668..81f107f0c9 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -53,6 +53,9 @@ # Get determinism _deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) +# CI test level +_TEST_LEVEL = os.getenv("NVTE_JAX_UNITTEST_LEVEL", "L0") + @pytest.fixture(autouse=True, scope="module") def init(): @@ -469,6 +472,24 @@ def _get_max_segments_per_sequence(self): return 1 def _check_configs(self): + # Trim SWA configs for L0 and L1 to reduce test time; need to trim more in future test refactoring. + if ( + self.window_size is not None + and (self.dropout_prob != 0.0 or self.attn_bias_type is not AttnBiasType.NO_BIAS) + ): + if _TEST_LEVEL == "L0" and ( + self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX + or self.dtype != jnp.bfloat16 + or self.attn_bias_type is not AttnBiasType.POST_SCALE_BIAS + or self.attn_mask_type is not AttnMaskType.NO_MASK + ): + pytest.skip("Trimmed SWA+bias/dropout config: only vanilla-softmax + bf16 + post_scale_bias + no-mask runs at L0") + if _TEST_LEVEL == "L1" and ( + self.dtype != jnp.float16 + or self.softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX + ): + pytest.skip("Trimmed SWA+bias/dropout config: only float16 + learnable-softmax runs at L1") + # TODO(rewang): probably adds this in is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") From 5b33337e10c1883e666a83e99716aa47618f3181 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:52:03 +0000 Subject: [PATCH 61/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_fused_attn.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 81f107f0c9..30f29c12d4 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -473,9 +473,8 @@ def _get_max_segments_per_sequence(self): def _check_configs(self): # Trim SWA configs for L0 and L1 to reduce test time; need to trim more in future test refactoring. - if ( - self.window_size is not None - and (self.dropout_prob != 0.0 or self.attn_bias_type is not AttnBiasType.NO_BIAS) + if self.window_size is not None and ( + self.dropout_prob != 0.0 or self.attn_bias_type is not AttnBiasType.NO_BIAS ): if _TEST_LEVEL == "L0" and ( self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX @@ -483,12 +482,16 @@ def _check_configs(self): or self.attn_bias_type is not AttnBiasType.POST_SCALE_BIAS or self.attn_mask_type is not AttnMaskType.NO_MASK ): - pytest.skip("Trimmed SWA+bias/dropout config: only vanilla-softmax + bf16 + post_scale_bias + no-mask runs at L0") + pytest.skip( + "Trimmed SWA+bias/dropout config: only vanilla-softmax + bf16 + post_scale_bias" + " + no-mask runs at L0" + ) if _TEST_LEVEL == "L1" and ( - self.dtype != jnp.float16 - or self.softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX + self.dtype != jnp.float16 or self.softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX ): - pytest.skip("Trimmed SWA+bias/dropout config: only float16 + learnable-softmax runs at L1") + pytest.skip( + "Trimmed SWA+bias/dropout config: only float16 + learnable-softmax runs at L1" + ) # TODO(rewang): probably adds this in is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): From 6461b66673adf0f3c27aa71642e0ddb8abb53f98 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 27 Jul 2026 08:50:49 -0700 Subject: [PATCH 62/63] fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 5 +- .../common/fused_attn/graph_cache_debug.h | 86 +++++++++---------- 2 files changed, 45 insertions(+), 46 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index bddaafb8d6..b3b9922abf 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -235,10 +235,9 @@ 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) { + if (message == nullptr) return; fused_attn_backend_message_buffer = std::move(reason); - if (message != nullptr) { - *message = fused_attn_backend_message_buffer.c_str(); - } + *message = fused_attn_backend_message_buffer.c_str(); } } // namespace diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 843e4156b4..523c0625cb 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -77,17 +78,14 @@ inline void print_counters(const char *event) { const EventCounters &f = counters(/*is_fwd=*/true); const EventCounters &b = counters(/*is_fwd=*/false); std::fprintf(stderr, - "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%llu exec=%llu hit=%llu miss=%llu | " - "bwd built=%llu exec=%llu hit=%llu miss=%llu\n", - event, thread_seq_id(), - static_cast(f.built.load(std::memory_order_relaxed)), - static_cast(f.exec.load(std::memory_order_relaxed)), - static_cast(f.hit.load(std::memory_order_relaxed)), - static_cast(f.miss.load(std::memory_order_relaxed)), - static_cast(b.built.load(std::memory_order_relaxed)), - static_cast(b.exec.load(std::memory_order_relaxed)), - static_cast(b.hit.load(std::memory_order_relaxed)), - static_cast(b.miss.load(std::memory_order_relaxed))); + "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%" PRIu64 " exec=%" PRIu64 + " hit=%" PRIu64 " miss=%" PRIu64 " | bwd built=%" PRIu64 " exec=%" PRIu64 + " hit=%" PRIu64 " miss=%" PRIu64 "\n", + event, thread_seq_id(), f.built.load(std::memory_order_relaxed), + f.exec.load(std::memory_order_relaxed), f.hit.load(std::memory_order_relaxed), + f.miss.load(std::memory_order_relaxed), b.built.load(std::memory_order_relaxed), + b.exec.load(std::memory_order_relaxed), b.hit.load(std::memory_order_relaxed), + b.miss.load(std::memory_order_relaxed)); std::fflush(stderr); } @@ -114,37 +112,39 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %-3s %-4s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%lld " - "bias=%lld wl=%lld wr=%lld brd=%d softmax=%lld scale_mode=%lld dropout=%g attn_scale=%g " - "qkv_dt=%lld o_dt=%lld do_dt=%lld dqkv_dt=%lld qkv_lay=%lld o_fmt=%lld do_fmt=%lld " - "dqkv_lay=%lld qkv_sif=%lld do_sif=%lld b=%lld h=%lld hg=%lld dqk=%lld dv=%lld sq=%lld " - "skv=%lld tq=%lld tkv=%lld bb=%lld btq=%lld btkv=%lld npk=%lld npv=%lld psk=%lld psv=%lld " - "mppk=%lld mppv=%lld bias_b=%lld bias_h=%lld bias_sq=%lld bias_skv=%lld\n", + "[FUSED-ATTN-CACHE] %-3s %-4s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d " + "mask=%" PRId64 " bias=%" PRId64 " wl=%" PRId64 " wr=%" PRId64 " brd=%d softmax=%" PRId64 + " scale_mode=%" PRId64 " dropout=%g attn_scale=%g qkv_dt=%" PRId64 " o_dt=%" PRId64 + " do_dt=%" PRId64 " dqkv_dt=%" PRId64 " qkv_lay=%" PRId64 " o_fmt=%" PRId64 " do_fmt=%" PRId64 + " dqkv_lay=%" PRId64 " qkv_sif=%" PRId64 " do_sif=%" PRId64 " b=%" PRId64 " h=%" PRId64 + " hg=%" PRId64 " dqk=%" PRId64 " dv=%" PRId64 " sq=%" PRId64 " skv=%" PRId64 " tq=%" PRId64 + " tkv=%" PRId64 " bb=%" PRId64 " btq=%" PRId64 " btkv=%" PRId64 " npk=%" PRId64 " npv=%" PRId64 + " psk=%" PRId64 " psv=%" PRId64 " mppk=%" PRId64 " mppv=%" PRId64 " bias_b=%" PRId64 + " bias_h=%" PRId64 " bias_sq=%" PRId64 " bias_skv=%" PRId64 "\n", pass, hit ? "HIT" : "MISS", thread_seq_id(), static_cast(c.is_training), static_cast(c.deterministic), static_cast(c.cuda_graph), static_cast(c.return_max_logit), static_cast(c.is_forward), - static_cast(c.attn_mask_type), static_cast(c.bias_type), - static_cast(c.window_size_left), static_cast(c.window_size_right), - static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), - static_cast(c.scaling_mode), static_cast(c.dropout), - static_cast(c.attn_scale), static_cast(c.qkv_dtype), - static_cast(c.o_dtype), static_cast(c.do_dtype), - static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), - static_cast(c.o_format), static_cast(c.do_format), - static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), - static_cast(c.do_scale_inv_format), static_cast(c.batch_size), - static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), - static_cast(c.head_dim_qk), static_cast(c.head_dim_v), - static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), - static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), - static_cast(c.bucketed_batch_size), - static_cast(c.bucketed_num_tokens_q), - static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), - static_cast(c.num_pages_v), static_cast(c.page_size_k), - static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), - static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), - static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), - static_cast(c.bias_seqlen_kv)); + static_cast(c.attn_mask_type), static_cast(c.bias_type), + static_cast(c.window_size_left), static_cast(c.window_size_right), + static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), + static_cast(c.scaling_mode), static_cast(c.dropout), + static_cast(c.attn_scale), static_cast(c.qkv_dtype), + static_cast(c.o_dtype), static_cast(c.do_dtype), + static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), + static_cast(c.o_format), static_cast(c.do_format), + static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), + static_cast(c.do_scale_inv_format), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.head_dim_qk), static_cast(c.head_dim_v), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), + static_cast(c.num_pages_v), static_cast(c.page_size_k), + static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), + static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), + static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), + static_cast(c.bias_seqlen_kv)); std::fflush(stderr); } @@ -215,10 +215,10 @@ inline void register_summary_once() { if (n == 0) continue; const double total_ms = static_cast(t.time_ns.load(std::memory_order_relaxed)) / 1e6; - std::fprintf( - stderr, - "[FUSED-ATTN-CACHE] %-3s %-22s | calls=%llu | time=%9.1f ms | avg=%9.3f ms/call\n", - pass, kStageNames[i], static_cast(n), total_ms, total_ms / n); + std::fprintf(stderr, + "[FUSED-ATTN-CACHE] %-3s %-22s | calls=%" PRIu64 + " | time=%9.1f ms | avg=%9.3f ms/call\n", + pass, kStageNames[i], n, total_ms, total_ms / n); } } std::fflush(stderr); From 00429e725abdc2d725aa4acd1a6db5dec031a175 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:08:41 +0000 Subject: [PATCH 63/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/fused_attn/graph_cache_debug.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 523c0625cb..033edf2ead 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -118,9 +118,9 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi " do_dt=%" PRId64 " dqkv_dt=%" PRId64 " qkv_lay=%" PRId64 " o_fmt=%" PRId64 " do_fmt=%" PRId64 " dqkv_lay=%" PRId64 " qkv_sif=%" PRId64 " do_sif=%" PRId64 " b=%" PRId64 " h=%" PRId64 " hg=%" PRId64 " dqk=%" PRId64 " dv=%" PRId64 " sq=%" PRId64 " skv=%" PRId64 " tq=%" PRId64 - " tkv=%" PRId64 " bb=%" PRId64 " btq=%" PRId64 " btkv=%" PRId64 " npk=%" PRId64 " npv=%" PRId64 - " psk=%" PRId64 " psv=%" PRId64 " mppk=%" PRId64 " mppv=%" PRId64 " bias_b=%" PRId64 - " bias_h=%" PRId64 " bias_sq=%" PRId64 " bias_skv=%" PRId64 "\n", + " tkv=%" PRId64 " bb=%" PRId64 " btq=%" PRId64 " btkv=%" PRId64 " npk=%" PRId64 + " npv=%" PRId64 " psk=%" PRId64 " psv=%" PRId64 " mppk=%" PRId64 " mppv=%" PRId64 + " bias_b=%" PRId64 " bias_h=%" PRId64 " bias_sq=%" PRId64 " bias_skv=%" PRId64 "\n", pass, hit ? "HIT" : "MISS", thread_seq_id(), static_cast(c.is_training), static_cast(c.deterministic), static_cast(c.cuda_graph), static_cast(c.return_max_logit), static_cast(c.is_forward),