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/88] 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/88] [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/88] 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/88] 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/88] [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/88] 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/88] [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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] [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/88] 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/88] [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/88] 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/88] [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/88] 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/88] [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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] 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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] 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/88] 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/88] 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/88] 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/88] 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/88] 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/88] 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/88] [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/88] 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/88] 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/88] 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/88] [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/88] 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/88] [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), From 988c9badad65e2533d6d73c94fde17beed3f69e2 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:46:26 -0700 Subject: [PATCH 64/88] fix torch.compile for get_backend Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/test_torch_compile.py | 2 +- .../attention/dot_product_attention/utils.py | 66 ++++++++++++------- 2 files changed, 43 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index d720ca7a4a..832876e825 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -645,7 +645,7 @@ def fn(x, params): monkeypatch.setattr( dpa_utils.tex, "get_fused_attn_backend", - lambda *args: dpa_utils.FusedAttnBackend["No_Backend"], + lambda *args: (dpa_utils.FusedAttnBackend["No_Backend"], "disabled by test"), ) def fn_no_backend(x, params): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 79377ea0c0..14f8d75a67 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -427,13 +427,24 @@ def error(self, *args, **kwargs): _no_op_logger = _NoOpLogger() +# torch.compile: +# sub-backend ids, i.e. the FusedAttnBackend values as plain ints, and the reverse mapping. +_FUSED_ATTN_BACKEND_IDS = { + name: int(member) for name, member in FusedAttnBackend.__members__.items() +} +_FUSED_ATTN_BACKENDS_BY_ID = {int(member): member for member in FusedAttnBackend} + + @torch.compiler.assume_constant_result -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).""" - fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) - return FusedAttnBackend.cast(fused_attention_backend), reject_message +def _get_fused_attn_backend(**fused_attn_kwargs): + """Constant-foldable tex.get_fused_attn_backend: the result depends only on the + attention config, so torch.compile can bake it in. Returns the sub-backend as a plain + int, i.e. one of `_FUSED_ATTN_BACKEND_IDS`, next to the rejection message. + """ + fused_attention_backend, reject_message = tex.get_fused_attn_backend( + FusedAttentionParams(**fused_attn_kwargs) + ) + return int(fused_attention_backend), reject_message def get_attention_backend( @@ -1505,7 +1516,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False # Filter: cuDNN support - fused_attention_backend = None + fused_attention_backend_id = None if use_fused_attention: 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" @@ -1617,13 +1628,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt 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) - - # NOTE: under torch.compile the numeric fields of fused_attn_params must not be + # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. - fused_attention_backend, reject_message = _get_fused_attn_backend(fused_attn_params) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: + fused_attention_backend_id, reject_message = _get_fused_attn_backend( + **fused_attn_kwargs + ) + if fused_attention_backend_id == _FUSED_ATTN_BACKEND_IDS["No_Backend"]: logger.debug( "Disabling FusedAttention: %s%s", reject_message, @@ -1634,21 +1645,21 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ), ) use_fused_attention = False - fused_attention_backend = None + fused_attention_backend_id = None break if ( use_fused_attention and has_score_mod - and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"] + and fused_attention_backend_id != _FUSED_ATTN_BACKEND_IDS["F16_arbitrary_seqlen"] ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " "F16/BF16 arbitrary-seqlen", - int(fused_attention_backend), + fused_attention_backend_id, ) use_fused_attention = False - fused_attention_backend = None + fused_attention_backend_id = None # Filter: Determinism # backend | deterministic # --------------------------------------------- @@ -1689,9 +1700,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt softmax_type, ) use_fused_attention = False - fused_attention_backend = None + fused_attention_backend_id = None if ( - fused_attention_backend == FusedAttnBackend["FP8"] + fused_attention_backend_id == _FUSED_ATTN_BACKEND_IDS["FP8"] and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1700,9 +1711,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt " < 9.19.0" ) use_fused_attention = False - fused_attention_backend = None + fused_attention_backend_id = None if ( - fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] + fused_attention_backend_id == _FUSED_ATTN_BACKEND_IDS["F16_arbitrary_seqlen"] and is_training and ( device_compute_capability < (9, 0) @@ -1712,7 +1723,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ): logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False - fused_attention_backend = None + fused_attention_backend_id = None # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 @@ -1785,8 +1796,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""), bool(available_backends[1]), ( - f" (sub-backend {int(fused_attention_backend)})" - if fused_attention_backend is not None + f" (sub-backend {fused_attention_backend_id})" + if fused_attention_backend_id is not None else "" ), bool(available_backends[2]), @@ -1810,11 +1821,18 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_flash_attention: selected_backend = f"FlashAttention ({str(flash_attention_backend)})" elif use_fused_attention: - selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" + selected_backend = f"FusedAttention (sub-backend {fused_attention_backend_id})" elif use_unfused_attention: selected_backend = "UnfusedDotProductAttention" logger.debug("Selected backend = %s.", selected_backend) + # Hand the sub-backend back as a FusedAttnBackend member. + fused_attention_backend = ( + None + if fused_attention_backend_id is None + else _FUSED_ATTN_BACKENDS_BY_ID[fused_attention_backend_id] + ) + return ( use_flash_attention, flash_attention_backend, From d374a3df0ea5e3e8b0571775379b93921d20171d Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:23:53 -0700 Subject: [PATCH 65/88] a cleaner way to make torch.compile work Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/test_torch_compile.py | 3 +- .../attention/dot_product_attention/utils.py | 59 ++++++++----------- .../pytorch/cpp_extensions/fused_attn.py | 24 ++------ 3 files changed, 30 insertions(+), 56 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 832876e825..3267953d27 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -40,6 +40,7 @@ from transformer_engine.pytorch.attention.dot_product_attention.backends import ( UnfusedDotProductAttention, ) +from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) @@ -645,7 +646,7 @@ def fn(x, params): monkeypatch.setattr( dpa_utils.tex, "get_fused_attn_backend", - lambda *args: (dpa_utils.FusedAttnBackend["No_Backend"], "disabled by test"), + lambda *args: (tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, "disabled by test"), ) def fn_no_backend(x, params): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 14f8d75a67..3c2c5a3b0c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -427,24 +427,20 @@ def error(self, *args, **kwargs): _no_op_logger = _NoOpLogger() -# torch.compile: -# sub-backend ids, i.e. the FusedAttnBackend values as plain ints, and the reverse mapping. -_FUSED_ATTN_BACKEND_IDS = { - name: int(member) for name, member in FusedAttnBackend.__members__.items() -} -_FUSED_ATTN_BACKENDS_BY_ID = {int(member): member for member in FusedAttnBackend} - - @torch.compiler.assume_constant_result def _get_fused_attn_backend(**fused_attn_kwargs): - """Constant-foldable tex.get_fused_attn_backend: the result depends only on the - attention config, so torch.compile can bake it in. Returns the sub-backend as a plain - int, i.e. one of `_FUSED_ATTN_BACKEND_IDS`, next to the rejection message. + """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). + + The config comes in as keyword arguments rather than as a FusedAttentionParams: + torch.compile materializes the arguments of a constant-folded call, and a + dataclass built by traced code arrives here with its fields reset to defaults. """ fused_attention_backend, reject_message = tex.get_fused_attn_backend( FusedAttentionParams(**fused_attn_kwargs) ) - return int(fused_attention_backend), reject_message + return FusedAttnBackend.cast(fused_attention_backend), reject_message def get_attention_backend( @@ -1516,7 +1512,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False # Filter: cuDNN support - fused_attention_backend_id = None + fused_attention_backend = None if use_fused_attention: 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" @@ -1631,10 +1627,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. - fused_attention_backend_id, reject_message = _get_fused_attn_backend( - **fused_attn_kwargs - ) - if fused_attention_backend_id == _FUSED_ATTN_BACKEND_IDS["No_Backend"]: + fused_attention_backend, reject_message = _get_fused_attn_backend(**fused_attn_kwargs) + if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug( "Disabling FusedAttention: %s%s", reject_message, @@ -1645,21 +1639,21 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ), ) use_fused_attention = False - fused_attention_backend_id = None + fused_attention_backend = None break if ( use_fused_attention and has_score_mod - and fused_attention_backend_id != _FUSED_ATTN_BACKEND_IDS["F16_arbitrary_seqlen"] + 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", - fused_attention_backend_id, + int(fused_attention_backend), ) use_fused_attention = False - fused_attention_backend_id = None + fused_attention_backend = None # Filter: Determinism # backend | deterministic # --------------------------------------------- @@ -1700,9 +1694,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt softmax_type, ) use_fused_attention = False - fused_attention_backend_id = None + fused_attention_backend = None if ( - fused_attention_backend_id == _FUSED_ATTN_BACKEND_IDS["FP8"] + fused_attention_backend == FusedAttnBackend["FP8"] and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1711,9 +1705,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt " < 9.19.0" ) use_fused_attention = False - fused_attention_backend_id = None + fused_attention_backend = None if ( - fused_attention_backend_id == _FUSED_ATTN_BACKEND_IDS["F16_arbitrary_seqlen"] + fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] and is_training and ( device_compute_capability < (9, 0) @@ -1723,7 +1717,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ): logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False - fused_attention_backend_id = None + fused_attention_backend = None # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 @@ -1796,8 +1790,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""), bool(available_backends[1]), ( - f" (sub-backend {fused_attention_backend_id})" - if fused_attention_backend_id is not None + f" (sub-backend {int(fused_attention_backend)})" + if fused_attention_backend is not None else "" ), bool(available_backends[2]), @@ -1821,18 +1815,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_flash_attention: selected_backend = f"FlashAttention ({str(flash_attention_backend)})" elif use_fused_attention: - selected_backend = f"FusedAttention (sub-backend {fused_attention_backend_id})" + selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" elif use_unfused_attention: selected_backend = "UnfusedDotProductAttention" logger.debug("Selected backend = %s.", selected_backend) - # Hand the sub-backend back as a FusedAttnBackend member. - fused_attention_backend = ( - None - if fused_attention_backend_id is None - else _FUSED_ATTN_BACKENDS_BY_ID[fused_attention_backend_id] - ) - return ( use_flash_attention, flash_attention_backend, diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 046019ee58..1d8955029e 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -110,7 +110,11 @@ class FusedAttnBackend(IntEnum): ``IntEnum`` is traceable by ``torch.compile``: comparisons constant-fold cleanly and instances safely cross the ``assume_constant_result`` boundary in ``get_attention_backend``. Lookup by name (``FusedAttnBackend["FP8"]``) - works the same way as with the dict this used to be. + works the same way as with the dict this used to be. Adding + ``__eq__``/``__ne__`` overrides is unnecessary (the inherited ``int`` + comparisons already match the pybind enum in both operand orders) and + harmful: a python ``__eq__`` would push dynamo from constant-folding the + comparison to inline-with-guard, breaking tracing. """ No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) @@ -131,24 +135,6 @@ def cast( return backend return cls(int(backend)) - def __eq__(self, other: object) -> bool: - # ``FusedAttnBackend`` is an ``IntEnum`` while ``NVTE_Fused_Attn_Backend`` - # is a pybind11 enum. Compare by integer value so the two enums stay - # equivalent regardless of the pybind11 version (the pybind ``__eq__`` - # handles the reverse order). - if isinstance(other, NVTE_Fused_Attn_Backend): - return int(self) == int(other) - return int.__eq__(self, other) - - def __ne__(self, other: object) -> bool: - result = self.__eq__(other) - if result is NotImplemented: - return result - return not result - - def __hash__(self) -> int: - return int.__hash__(self) - # Fail fast at import time if a new enumerator is added on the C++ side # without being mirrored above. From 288046c803bd932818c12fdaeeec60ebd512bdd0 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:27:32 -0700 Subject: [PATCH 66/88] skip fused attn checks for flash tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention_with_cp.py | 1 + tests/pytorch/utils.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 9ffdb865fa..c9d19588f7 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -380,6 +380,7 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type qkv_layout="_".join([qkv_format] * 3), cp_size=num_gpus, cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, + skip_fused_attn=True, ) flash_attn_supported, *_ = available_backends if not flash_attn_supported: diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index cdab36b2c8..90fbcc16b5 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -351,11 +351,18 @@ def get_available_attention_backends( score_mod_bprop: bool = False, cp_size: int = 1, cp_size_a2a: int = 1, + skip_fused_attn: bool = False, ) -> Tuple[List, List]: - """Check for all available attention backends that support a model configuration""" + """Check for all available attention backends that support a model configuration + + Set `skip_fused_attn=True` to leave fused attention out of the query. The reported + fused-attention backends are then empty, while the FlashAttention and unfused results + are unaffected. This skips cuDNN's support checks, which build and cache a graph per + configuration. + """ os.environ["NVTE_FLASH_ATTN"] = "1" - os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_FUSED_ATTN"] = "0" if skip_fused_attn else "1" os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True alibi_slopes_shape = None From 1d10f80f87a501f256852d9bfac17efddc8a3f75 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:28:09 -0700 Subject: [PATCH 67/88] add tq/tkv to per-step cp configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../dot_product_attention/context_parallel.py | 33 +++++++++++++++---- .../attention/dot_product_attention/utils.py | 4 +++ 2 files changed, 31 insertions(+), 6 deletions(-) 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 0c611051f7..1165a98387 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4925,6 +4925,8 @@ def cp_per_step_configs( *, max_seqlen_q, max_seqlen_kv, + num_tokens_q, + num_tokens_kv, num_heads, num_gqa_groups, attn_mask_type, @@ -4941,11 +4943,13 @@ def cp_per_step_configs( 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): + def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): return { "attn_mask_type": mask, "max_seqlen_q": s_q, "max_seqlen_kv": s_kv, + "num_tokens_q": t_q, + "num_tokens_kv": t_kv, "num_attn_heads": heads, "num_gqa_groups": gqa, "window_size_left": window_left, @@ -4963,6 +4967,8 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right): num_heads // cp_size, num_gqa_groups // cp_size, bottom_right_diagonal, + num_tokens_q * cp_size, + num_tokens_kv * cp_size, ) ] @@ -4973,10 +4979,20 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right): 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 + t_q = num_tokens_q // 2 # 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) + config( + mask, + s_q, + s_kv, + num_heads, + num_gqa_groups, + br, + t_q, + num_tokens_kv * cp_size * s_kv // max_seqlen_kv if max_seqlen_kv else 0, + ) for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) ] @@ -4986,15 +5002,20 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right): gqa = num_gqa_groups // cp_size_a2a r_q = max_seqlen_q // p2p_size r_kv = max_seqlen_kv // p2p_size + # The tensors handed to this rank already correspond to (r_q, r_kv), so the token counts + # need no rescaling here; they only follow the halving below. + t_q, t_kv = num_tokens_q, num_tokens_kv 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, t_q, t_kv)] 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 + attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv + ), # diagonal + config( + padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal, t_q, t_kv // 2 ), # lower-triangle config( - padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal + padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal, t_q // 2, t_kv ), # 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 3c2c5a3b0c..8cbf6342c4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1596,6 +1596,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cp_size_a2a, 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_heads=num_heads, num_gqa_groups=num_gqa_groups, attn_mask_type=attn_mask_type, @@ -1614,6 +1616,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt attn_mask_type=AttnMaskType[step_config["attn_mask_type"]], max_seqlen_q=step_seqlen_q, max_seqlen_kv=step_seqlen_kv, + num_tokens_q=step_config["num_tokens_q"], + num_tokens_kv=step_config["num_tokens_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"], From e3300ea22c5b57fb85ac042342a1e5dfcf00f186 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:30:15 +0000 Subject: [PATCH 68/88] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/dot_product_attention/context_parallel.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 1165a98387..f11f52d68e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -5008,9 +5008,7 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): if not is_causal: return [config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv)] return [ - config( - attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv - ), # diagonal + config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv), # diagonal config( padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal, t_q, t_kv // 2 ), # lower-triangle From 71e898266e3c2ca2f4a716b065781069cc3ec0a8 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:55:20 -0700 Subject: [PATCH 69/88] fix jax CI Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../jax/attention_context_parallel.py | 31 ++++++++++--------- docs/examples/jax/test_attention.py | 31 ++++++++++--------- 2 files changed, 32 insertions(+), 30 deletions(-) diff --git a/docs/examples/jax/attention_context_parallel.py b/docs/examples/jax/attention_context_parallel.py index 1557a30b7c..982ecc11b2 100644 --- a/docs/examples/jax/attention_context_parallel.py +++ b/docs/examples/jax/attention_context_parallel.py @@ -245,21 +245,22 @@ def context_parallel_supported() -> Tuple[bool, str]: return False, f"needs {cp_size} GPUs" has_kernel = is_fused_attn_kernel_available( - True, - dtype, - dtype, - QKVLayout.THD_THD_THD, - AttnBiasType.NO_BIAS, - AttnMaskType.PADDING_CAUSAL_MASK, - AttnSoftmaxType.VANILLA_SOFTMAX, - 0.0, - num_query_heads, - num_kv_heads, - seq, - seq, - head_dim, - head_dim, - window_size, + is_training=True, + batch_size=batch, + q_dtype=dtype, + kv_dtype=dtype, + qkv_layout=QKVLayout.THD_THD_THD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=num_query_heads, + kv_num_heads=num_kv_heads, + q_max_seqlen=seq, + kv_max_seqlen=seq, + head_dim_qk=head_dim, + head_dim_v=head_dim, + window_size=window_size, ) if not has_kernel: return False, "no fused attention kernel for the THD SWA shape" diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py index 3cc08271dc..b5db8c3535 100644 --- a/docs/examples/jax/test_attention.py +++ b/docs/examples/jax/test_attention.py @@ -87,21 +87,22 @@ def _context_parallel_supported(): return False, f"needs {cp_size} GPUs" has_kernel = is_fused_attn_kernel_available( - True, - jnp.bfloat16, - jnp.bfloat16, - QKVLayout.THD_THD_THD, - AttnBiasType.NO_BIAS, - AttnMaskType.PADDING_CAUSAL_MASK, - AttnSoftmaxType.VANILLA_SOFTMAX, - 0.0, - 128, - 8, - 65536, - 65536, - 128, - 128, - (8192, 0), + is_training=True, + batch_size=2, + q_dtype=jnp.bfloat16, + kv_dtype=jnp.bfloat16, + qkv_layout=QKVLayout.THD_THD_THD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=128, + kv_num_heads=8, + q_max_seqlen=65536, + kv_max_seqlen=65536, + head_dim_qk=128, + head_dim_v=128, + window_size=(8192, 0), ) if not has_kernel: return False, "no fused attention kernel for the THD SWA shape" From 8fdd81d291a96ef9aebc0136894b90cb040eed0c Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:22:08 -0700 Subject: [PATCH 70/88] temporary changes: cache debug, timers, single flight, is_probe, dry-run, still build plans in probes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 236 +++++----- tests/pytorch/test_torch_compile.py | 1 + tests/pytorch/utils.py | 63 ++- .../common/fused_attn/config_and_params.cpp | 89 +++- .../common/fused_attn/config_and_params.h | 17 +- .../common/fused_attn/fused_attn.cpp | 28 +- .../fused_attn_f16_arbitrary_seqlen.cu | 90 ++-- .../common/fused_attn/fused_attn_fp8.cu | 90 ++-- .../common/fused_attn/graph_cache_debug.h | 428 ++++++++++++++++-- .../include/transformer_engine/fused_attn.h | 15 +- .../dot_product_attention/__init__.py | 16 +- .../dot_product_attention/context_parallel.py | 25 +- .../dot_product_attention.py | 140 +++++- .../attention/dot_product_attention/utils.py | 33 +- 14 files changed, 998 insertions(+), 273 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index ca93fecc9e..2b9b64a026 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -29,6 +29,7 @@ _attention_backends, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + FUSED_ATTN_BWD_REJECT_PREFIX, FlashAttentionUtils, check_set_window_size, ) @@ -62,6 +63,7 @@ ModelConfig, dtype_tols, get_available_attention_backends, + probe_attention_backends, ) # Check if hardware supports FP8 attention. @@ -103,18 +105,18 @@ def reset_global_fp8_state(): model_configs_base = { # test: ModelConfig(b, sq, hq, dqk) - "base_1_0": ModelConfig(8, 128, 16, 64), + "base_1_0": ModelConfig(8, 131072, 16, 64), "base_1_1": ModelConfig(4, 128, 16, 64, max_seqlen_kv=256), "base_2_0": ModelConfig(2, 2048, 24, 128), "base_2_1": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096), - "base_3_0": ModelConfig(8, 1, 16, 128, max_seqlen_kv=2048), - "base_3_1": ModelConfig(8, 1, 16, 256, max_seqlen_kv=2048), - "base_4_0": ModelConfig(8, 1, 16, 192, max_seqlen_kv=2048), - "base_4_1": ModelConfig(8, 128, 16, 192, max_seqlen_kv=2048), - "base_5_0": ModelConfig(8, 1, 16, 512, max_seqlen_kv=2048), - "base_5_1": ModelConfig(8, 128, 16, 512, max_seqlen_kv=2048), - "base_6_0": ModelConfig(8, 1, 16, 1024, max_seqlen_kv=2048), - "base_6_1": ModelConfig(8, 128, 16, 1024, max_seqlen_kv=2048), + # "base_3_0": ModelConfig(8, 1, 16, 128, max_seqlen_kv=2048), + # "base_3_1": ModelConfig(8, 1, 16, 256, max_seqlen_kv=2048), + # "base_4_0": ModelConfig(8, 1, 16, 192, max_seqlen_kv=2048), + # "base_4_1": ModelConfig(8, 128, 16, 192, max_seqlen_kv=2048), + # "base_5_0": ModelConfig(8, 1, 16, 512, max_seqlen_kv=2048), + # "base_5_1": ModelConfig(8, 128, 16, 512, max_seqlen_kv=2048), + # "base_6_0": ModelConfig(8, 1, 16, 1024, max_seqlen_kv=2048), + # "base_6_1": ModelConfig(8, 128, 16, 1024, max_seqlen_kv=2048), } @@ -178,17 +180,20 @@ def test_dot_product_attention( "Setting is_training to False as cuDNN does not support dbias for" f" {config.bias_shape=} " ) - available_backends, _, fused_attn_backends = get_available_attention_backends( + available_backends, _, fused_attn_backends, reject_reason = get_available_attention_backends( config, qkv_dtype=dtype, qkv_layout=qkv_layout, pad_between_seqs=pad_between_seqs, is_training=is_training, deterministic=_deterministic, + return_reason=True, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if not fused_attn_supported: + # Retry in inference mode only when the backward pass alone is what fused attention does not + # support; for any other reason dropping is_training cannot make it available. + if not fused_attn_supported and (reject_reason or "").startswith(FUSED_ATTN_BWD_REJECT_PREFIX): is_training = False available_backends, _, fused_attn_backends = get_available_attention_backends( config, @@ -205,6 +210,7 @@ def test_dot_product_attention( pytest.skip("Less than two backends to compare.") # UnfusedDotProductAttention backend + unfused_attn_supported=False if unfused_attn_supported: unfused_attn_fwd, unfused_max_logit, unfused_attn_bwd = _run_dot_product_attention( dtype, @@ -230,6 +236,7 @@ def test_dot_product_attention( ) # FlashAttention backend + flash_attn_supported = False if flash_attn_supported: flash_attn_fwd, _, flash_attn_bwd = _run_dot_product_attention( dtype, @@ -1584,35 +1591,40 @@ def test_transformer_layer( config = model_configs[model] tols = dict(atol=5e-2, rtol=5e-2) - # Test backend availability + # Test backend availability. Dry-run the module under test so the query uses the exact + # configuration it resolves, rather than a restatement of it that can drift. A decoder + # layer runs self-attention and then cross-attention, and needs both to be supported. is_training = True - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, - qkv_dtype=dtype, - qkv_layout=( - qkv_format.replace("hd", "h3d") if fused_qkv_params else qkv_format.replace("hd", "3hd") - ), - is_training=is_training, - deterministic=_deterministic, - ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if not fused_attn_supported: - is_training = False - available_backends, _, fused_attn_backends = get_available_attention_backends( + num_attn_sites = 2 if config.attn_type == "cross" else 1 + + def probe(is_training): + return probe_attention_backends( + _run_transformer_layer, + dtype, config, - qkv_dtype=dtype, - qkv_layout=( - qkv_format.replace("hd", "h3d") - if fused_qkv_params - else qkv_format.replace("hd", "3hd") - ), - is_training=is_training, - deterministic=_deterministic, + "", + ckpt_attn, + qkv_format, + fused_qkv_params, + RoPE, + is_training, + num_attn_sites=num_attn_sites, ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends + + dry_run = probe(is_training) + # Retry in inference mode only when the backward pass alone is unsupported (see + # test_dot_product_attention). + if not dry_run.fused_supported and (dry_run.fused_attn_reject_reason or "").startswith( + FUSED_ATTN_BWD_REJECT_PREFIX + ): + is_training = False + dry_run = probe(is_training) + flash_attn_supported = dry_run.flash_supported + fused_attn_supported = dry_run.fused_supported + unfused_attn_supported = dry_run.unfused_supported # Skip if only unfused backend is supported - if (len(fused_attn_backends) + flash_attn_supported + unfused_attn_supported) < 2: + if (fused_attn_supported + flash_attn_supported + unfused_attn_supported) < 2: pytest.skip("Less than two backends to compare.") # Skip if qkv_format = thd and "padding" not in attn_mask_type if qkv_format == "thd" and "padding" not in config.attn_mask_type: @@ -1645,6 +1657,7 @@ def test_transformer_layer( ) # FlashAttention backend + flash_attn_supported = False if flash_attn_supported: flash_attn_fwd, flash_attn_bwd = _run_transformer_layer( dtype, @@ -1727,17 +1740,13 @@ def _run_transformer_layer( ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Run TransformerLayer module with one forward pass and one backward pass""" - # Set RNG and environment variables + # Set RNG and environment variables. An empty `backend` leaves the caller's choice of + # enabled backends alone, so probe_attention_backends() can ask about all of them. reset_rng_states() - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - if backend == "FlashAttention": - os.environ["NVTE_FLASH_ATTN"] = "1" - if backend == "FusedAttention": - os.environ["NVTE_FUSED_ATTN"] = "1" - if backend == "UnfusedDotProductAttention": - os.environ["NVTE_UNFUSED_ATTN"] = "1" + if backend: + os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FlashAttention" else "0" + os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "FusedAttention" else "0" + os.environ["NVTE_UNFUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention" else "0" _attention_backends["backend_selection_requires_update"] = True # Create input tensor @@ -2073,49 +2082,50 @@ def get_model(dtype, config): attn_mask_type = "causal" model_configs_fp8_vs_f16 = { # test: ModelConfig(b, sq, hq, dqk) - "fp8_9": ModelConfig( - 2, - 2048, - 128, - 192, - head_dim_v=128, - ), - "fp8_10": ModelConfig( - 2, - 2048, - 128, - 192, - head_dim_v=128, - attn_mask_type="causal", - ), - "fp8_11": ModelConfig( - 2, - 2048, - 128, - 192, - head_dim_v=128, - attn_mask_type="causal_bottom_right", - ), - "fp8_12": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), - "fp8_13": ModelConfig( - 2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(128, 0) - ), - "fp8_14": ModelConfig(2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), - "fp8_15": ModelConfig(1, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0)), - "fp8_16": ModelConfig( - 1, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" - ), - "fp8_17": ModelConfig( - 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" - ), - "fp8_18": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), - "fp8_19": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), - "fp8_20": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), + # "fp8_9": ModelConfig( + # 2, + # 2048, + # 128, + # 192, + # head_dim_v=128, + # ), + # "fp8_10": ModelConfig( + # 2, + # 2048, + # 128, + # 192, + # head_dim_v=128, + # attn_mask_type="causal", + # ), + # "fp8_11": ModelConfig( + # 2, + # 2048, + # 128, + # 192, + # head_dim_v=128, + # attn_mask_type="causal_bottom_right", + # ), + # "fp8_12": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), + # "fp8_13": ModelConfig( + # 2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(128, 0) + # ), + # "fp8_14": ModelConfig(2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), + "fp8_15": ModelConfig(1, 8192, 64, 64, #attn_mask_type="causal", #window_size=(128, 0) + ), + # "fp8_16": ModelConfig( + # 1, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" + # ), + # "fp8_17": ModelConfig( + # 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" + # ), + # "fp8_18": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), + # "fp8_19": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), + # "fp8_20": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), } -param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16] -qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"] -qkv_format_fp8_vs_f16 = ["bshd", "sbhd"] +param_types_fp8_vs_f16 = [torch.bfloat16] #[torch.float16, torch.bfloat16] +qkv_layout_fp8_vs_f16 = ["sbhd_sbhd_sbhd"] #["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"] +qkv_format_fp8_vs_f16 = ["sbhd"] #["bshd", "sbhd"] @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") @@ -2123,11 +2133,11 @@ def get_model(dtype, config): @pytest.mark.parametrize("dtype", param_types_fp8_vs_f16) @pytest.mark.parametrize("model", model_configs_fp8_vs_f16.keys()) @pytest.mark.parametrize("qkv_format", qkv_format_fp8_vs_f16) -@pytest.mark.parametrize("input_layernorm", [True, False]) -@pytest.mark.parametrize("fp8_dpa_bwd", [True, False]) -@pytest.mark.parametrize("RoPE", [True, False]) -@pytest.mark.parametrize("is_training", [True, False]) -@pytest.mark.parametrize("scaling_mode", ["delayed", "current", "mxfp8"]) +@pytest.mark.parametrize("input_layernorm", [False]) #True, False]) +@pytest.mark.parametrize("fp8_dpa_bwd", [True])#, False]) +@pytest.mark.parametrize("RoPE", [False]) #True, False]) +@pytest.mark.parametrize("is_training", [True]) #, False]) +@pytest.mark.parametrize("scaling_mode", ["delayed"]) #, "current", "mxfp8"]) def test_mha_fp8_vs_f16( dtype, model, @@ -2169,29 +2179,37 @@ def test_mha_fp8_vs_f16( ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe - available_backends, _, _ = get_available_attention_backends( + # Dry-run the module under test so the query uses the exact configuration it resolves, + # rather than a restatement of it that can drift. + fp8_probe = probe_attention_backends( + _run_mha_fp8_vs_f16, + dtype, config, - qkv_dtype=torch.float8_e4m3fn, - nominal_dtype=dtype, - qkv_layout=qkv_format.replace("hd", "h3d"), - fp8=True, - fp8_meta=fp8_meta, - is_training=is_training, - deterministic=_deterministic, + True, + qkv_format, + input_layernorm, + RoPE, + is_training, + fp8_recipe, ) - flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends - available_backends, _, fused_attn_backends = get_available_attention_backends( + flash_attn_supported = fp8_probe.flash_supported + fused_attn_supported_fp8 = fp8_probe.fused_supported + f16_probe = probe_attention_backends( + _run_mha_fp8_vs_f16, + dtype, config, - qkv_dtype=dtype, - qkv_layout=qkv_format.replace("hd", "h3d"), - is_training=is_training, - deterministic=_deterministic, + False, + qkv_format, + input_layernorm, + RoPE, + is_training, + fp8_recipe, ) - _, fused_attn_supported_f16, _ = available_backends + fused_attn_supported_f16 = f16_probe.fused_supported if flash_attn_supported + fused_attn_supported_fp8 < 1: - pytest.skip("No FP8 attention backend available.") + pytest.skip(fp8_probe.fused_attn_reject_reason or "No FP8 attention backend available.") if not fused_attn_supported_f16: - pytest.skip("No reference backend available.") + pytest.skip(f16_probe.fused_attn_reject_reason or "No reference backend available.") if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" @@ -2300,6 +2318,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attention_type="self", qkv_weight_interleaved=True, qkv_format=qkv_format, + window_size=config.window_size, + softmax_type=config.softmax_type, ).to(dtype=dtype, device="cuda") if not is_training: mha = mha.eval() diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index dc0da5106e..05f8a36c4a 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -583,6 +583,7 @@ def fn(x, params): fused_attention_backend, use_unfused_attention, _, + _, ) = dpa_utils.get_attention_backend(params) # Encode the full selection (enabled backends + fused sub-backend) in # the tensor value: without a tensor op dynamo skips the frame entirely diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 90fbcc16b5..4652d433e7 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -20,7 +20,11 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch import InferenceParams, QuantizedTensor from transformer_engine.pytorch import DType -from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends +from transformer_engine.pytorch.attention.dot_product_attention import ( + DryRunResult, + dry_run_backend_selection, + _attention_backends, +) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, AttentionParams, @@ -336,6 +340,45 @@ def logging_context(highest_level=logging.WARNING): logging.disable(previous_level) +def probe_attention_backends(run_fn, *args, num_attn_sites: int = 1, **kwargs) -> DryRunResult: + """Which backends support the configuration a module actually produces. + + `run_fn(*args, **kwargs)` should build and call the module under test exactly as a real + run would. It is aborted inside `DotProductAttention`, once backend selection is known + but before any attention executes. + + Prefer this over `get_available_attention_backends()` for module-level tests. The + latter needs the caller to restate the configuration, which means predicting what the + module derives internally, e.g. the `qkv_layout` that `MultiheadAttention` gets from + its packed projection output. Whenever such a prediction drifts from the module, cuDNN + builds a graph under a cache key that is never executed. + + All three backends are enabled before `run_fn` is called, so a run function that does + not force a backend itself is probed against all of them. A run function that does + force one, e.g. `_run_transformer_layer(backend=...)`, is instead probed under exactly + the environment it will really use; read the matching `*_supported` property. + + Set `num_attn_sites` above 1 for a module that reaches `DotProductAttention` more than + once, e.g. a `TransformerLayer` with `layer_type="decoder"`, which runs self-attention + and then cross-attention. The `*_supported` properties then require every site to be + supported, which is what the module needs to run. + """ + os.environ["NVTE_FLASH_ATTN"] = "1" + os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + + with dry_run_backend_selection(stop_after=num_attn_sites) as dry_run: + run_fn(*args, **kwargs) + + assert len(dry_run.probes) == num_attn_sites, ( + f"dry run reached {len(dry_run.probes)} attention site(s), expected" + f" {num_attn_sites}; check num_attn_sites" + ) + _attention_backends["backend_selection_requires_update"] = True + return dry_run + + def get_available_attention_backends( config: ModelConfig, qkv_dtype: torch.dtype, @@ -352,6 +395,7 @@ def get_available_attention_backends( cp_size: int = 1, cp_size_a2a: int = 1, skip_fused_attn: bool = False, + return_reason: bool = False, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration @@ -359,6 +403,11 @@ def get_available_attention_backends( fused-attention backends are then empty, while the FlashAttention and unfused results are unaffected. This skips cuDNN's support checks, which build and cache a graph per configuration. + + Set `return_reason=True` to append the fused-attention rejection reason to the returned + tuple. A reason starting with `FUSED_ATTN_BWD_REJECT_PREFIX` means only the backward pass + is unsupported, i.e. re-querying with `is_training=False` may report fused attention as + available; any other reason means it will not. """ os.environ["NVTE_FLASH_ATTN"] = "1" @@ -442,6 +491,7 @@ def test(): fused_attention_backend, use_unfused_attention, available_backends, + fused_attention_reject_reason, ) = get_attention_backend(attention_params) # Check if FA3 is an available backend when num_splits != 1 if available_backends[0]: @@ -455,16 +505,23 @@ def test(): _attention_backends["fused_attention_backend"] = fused_attention_backend _attention_backends["use_unfused_attention"] = use_unfused_attention _attention_backends["backend_selection_requires_update"] = False - return available_backends, flash_attention_backend, fused_attention_backend + return ( + available_backends, + flash_attention_backend, + fused_attention_backend, + fused_attention_reject_reason, + ) backends = {1: "F16_arbitrary_seqlen", 2: "FP8"} if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() _attention_backends["backend_selection_requires_update"] = True - available_backends, flash_attention_backend, fused_attention_backend = test() + available_backends, flash_attention_backend, fused_attention_backend, reject_reason = test() if fused_attention_backend in (FusedAttnBackend[name] for name in backends.values()): fused_attn_backends.append(fused_attention_backend) + if return_reason: + return available_backends, flash_attention_backend, fused_attn_backends, reject_reason return available_backends, flash_attention_backend, 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 ca4214dac3..358beb7641 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -9,10 +9,13 @@ #include #include +#include +#include #include #include "../common.h" #include "../util/cuda_runtime.h" +#include "graph_cache_debug.h" namespace { @@ -24,6 +27,69 @@ void uint8_to_bool(const void *in, bool &out) { out = static_cast(*reinterpret_cast(in)); } +// Whether all visible CUDA devices can execute each other's cuDNN graphs. A plan is compiled +// against both the SM architecture and the SM count, and cuDNN requires both to match for a +// graph built on one device to run on another, so a difference in either one rules out sharing. +// SM counts differ across devices in practice even at a fixed arch, e.g. when MIG partitions +// or a harvested SKU are mixed in. +bool all_visible_devices_share_plans() { + const int n = transformer_engine::cuda::num_devices(); + if (n <= 1) return true; + const int arch0 = transformer_engine::cuda::sm_arch(0); + const int sm_count0 = transformer_engine::cuda::sm_count(0); + for (int i = 1; i < n; ++i) { + if (transformer_engine::cuda::sm_arch(i) != arch0) return false; + if (transformer_engine::cuda::sm_count(i) != sm_count0) return false; + } + return true; +} + +// Width reserved for the SM count in the packed cache key below. +constexpr int kSmCountBits = 16; + +// Scope of the fused-attention graph cache across devices in a single process. +// >= 0 : shared key packing (SM arch, SM count) -- all devices reuse one graph per shape +// (homogeneous node; a cuDNN plan compiled for this arch and SM count is valid on +// every device). +// -1 : per-device -- key by device id (heterogeneous node, or forced off). +// Computed once. The two schemes are never mixed within a process, so the packed values need +// not avoid the small device ids. +// +// NVTE_FUSED_ATTN_CACHE_PER_DEVICE=1 forces per-device keying. It is a debug-only +// escape hatch, not a tuning knob: the homogeneity check above is what keeps a +// mixed-arch node correct, so the only reasons to set it are A/B comparison +// against the old behavior, or working around a cuDNN plan-portability bug in the +// field without a rebuild. +int fused_attn_cache_arch_key() { + static const int key = [] () -> int { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_PER_DEVICE"); + const bool force_per_device = (e != nullptr && e[0] != '\0' && e[0] != '0'); + const int n = transformer_engine::cuda::num_devices(); + const bool per_device = force_per_device || !all_visible_devices_share_plans(); + const int arch = per_device ? -1 : transformer_engine::cuda::sm_arch(0); + const int sm_count = per_device ? -1 : transformer_engine::cuda::sm_count(0); + const int result = per_device ? -1 : ((arch << kSmCountBits) | sm_count); + // One-shot: make the resolved cache scope unambiguous in the diagnostics. + if (transformer_engine::fused_attn::graph_cache_debug::enabled()) { + const char *tag = transformer_engine::fused_attn::graph_cache_debug::process_tag().c_str(); + if (result >= 0) { + std::fprintf(stderr, + "\n[FUSED-ATTN-CACHE] %s | cache scope = arch %d + %d SM(s) (key %d) shared " + "across %d device(s)\n", + tag, arch, sm_count, result, n); + } else { + std::fprintf(stderr, + "\n[FUSED-ATTN-CACHE] %s | cache scope = per-device across %d device(s) (%s)\n", + tag, n, force_per_device ? "forced by NVTE_FUSED_ATTN_CACHE_PER_DEVICE" + : "devices differ in arch or SM count"); + } + std::fflush(stderr); + } + return result; + }(); + return key; +} + } // namespace namespace transformer_engine { @@ -96,8 +162,14 @@ 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(); + // Scope the graph cache across devices in a single process. cuDNN plans are compiled against + // a specific SM architecture and SM count, so on a node whose devices agree on both the plan + // built on one device is valid on all of them: key by (arch, SM count) so a shape is built + // once and shared (single-flight collapses the otherwise-per-device duplicate builds). + // When the devices differ in either one (or keying is forced off) fall back to per-device + // keying, since a plan is not portable in that case. + const int arch_key = fused_attn_cache_arch_key(); + cache_cfg.device_id = (arch_key >= 0) ? arch_key : cuda::current_device(); // Normalize bottom_right_diagonal const bool has_window = cache_cfg.window_size_left != -1 || cache_cfg.window_size_right != -1; @@ -121,7 +193,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { } cache_cfg.num_tokens_q = 0; cache_cfg.num_tokens_kv = 0; - const bool bucket_batch = !is_forward || !cache_cfg.uses_cu_seqlens_directly; + const bool bucket_batch = !check_forward || !cache_cfg.uses_cu_seqlens_directly; if (bucket_batch) { cache_cfg.batch_size = cache_cfg.bucketed_batch_size; } @@ -133,7 +205,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { // 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 (check_forward) { cache_cfg.do_dtype = kNVTEBFloat16; cache_cfg.dqkv_dtype = kNVTEBFloat16; cache_cfg.do_format = NVTE_QKV_Format_NOT_SET; @@ -150,7 +222,10 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig FusedAttnFwdParams::make_config() const { const FusedAttnFwdParams ¶ms = *this; FusedAttnConfig cfg{}; - cfg.is_forward = true; + // Forward execution: only the forward graph is run, so do not pay for a backward support + // check whose graph this call will never execute. + cfg.check_forward = true; + cfg.check_backward = false; cfg.is_training = params.is_training; cfg.deterministic = false; cfg.cuda_graph = params.cuda_graph; @@ -255,6 +330,10 @@ FusedAttnConfig FusedAttnFwdParams::make_config() const { FusedAttnConfig FusedAttnBwdParams::make_config() const { const FusedAttnBwdParams ¶ms = *this; FusedAttnConfig cfg{}; + // Backward execution: only the backward graph is run. check_forward=false also selects the + // backward key normalization in make_cache_key(). + cfg.check_forward = false; + cfg.check_backward = true; cfg.is_training = true; cfg.deterministic = params.deterministic; 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 ebc5b3eb07..1218c7860b 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -79,11 +79,18 @@ struct FusedAttnConfig { 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. - - // Direction to build the cuDNN graph for; steers make_cache_key() normalization. - bool is_forward = false; + // Filled by derive() or set by caller (i.e. check_forward). Added for convinence purposes and do + // not represent any graph properties. + + // Which directions nvte_get_fused_attn_backend_v2() runs a support check for. The execution + // entry points each run one direction and ask about that one only; a support query leaves the + // defaults and asks about both. Checking a direction that is never executed builds a cuDNN + // graph under a cache key nothing consumes. + // + // check_forward doubles as the direction to build the cuDNN graph for, steering + // make_cache_key() normalization, so it must stay true wherever a forward graph is built. + bool check_forward = true; + bool check_backward = true; // 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; diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index b3b9922abf..ebabe79844 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -300,15 +300,17 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi std::to_string(static_cast(qkv_format)) + "."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - 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 (cfg.check_forward) { + 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 (cfg.is_training && !cfg.is_forward) { + if (cfg.is_training && cfg.check_backward) { std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); if (!bwd_reason.empty()) { - set_message(message, std::move(bwd_reason)); + set_message(message, NVTE_FUSED_ATTN_BWD_REJECT_PREFIX + std::move(bwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } @@ -325,15 +327,17 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - std::string fwd_reason = is_supported_f16_fwd(cfg, handle); - if (!fwd_reason.empty()) { - set_message(message, std::move(fwd_reason)); - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (cfg.check_forward) { + std::string fwd_reason = is_supported_f16_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } } - if (cfg.is_training && !cfg.is_forward) { + if (cfg.is_training && cfg.check_backward) { std::string bwd_reason = is_supported_f16_bwd(cfg, handle); if (!bwd_reason.empty()) { - set_message(message, std::move(bwd_reason)); + set_message(message, NVTE_FUSED_ATTN_BWD_REJECT_PREFIX + std::move(bwd_reason)); 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 b7c7a349af..22ac9481e2 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 @@ -155,23 +155,32 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // 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; + static graph_cache::SingleFlight sdpa_f16_fprop_cache_sf; - // Get plan from cache if cache is available, otherwise create one + // Get plan from cache if available; otherwise build it exactly once across + // threads (single-flight), so concurrent misses of the same key don't each + // compile and discard an identical graph. auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - // Lock the map lookup, not the build, so different graphs can build in parallel - graph_and_tensors cached_graph{}; - bool cache_hit = false; + auto &sf = sdpa_f16_fprop_cache_sf; { - std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); + std::unique_lock lock(sf.mutex); + // Wait until the graph is cached, or no other thread is building this key. + sf.cv.wait(lock, [&] { + return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; + }); auto it = cache.find(descriptor); - 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; + if (it != cache.end()) { + graph_and_tensors cached_graph = it->second; // copy under the lock + lock.unlock(); + graph_cache_debug::record_cache_lookup("fwd", /*hit=*/true, cfg, descriptor.device_id); + return cached_graph; + } + // Claim the build for this key, so a concurrent miss waits instead of + // compiling an identical graph. No claim means no waiting. + if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); } + graph_cache_debug::record_cache_lookup("fwd", /*hit=*/false, cfg, descriptor.device_id); + graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -454,10 +463,11 @@ 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); 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. + // Insert our graph. With single-flight we are normally the only builder + // for this key; insert() still tolerates a pre-existing entry and returns + // it. claim_guard releases the build claim and wakes waiters on return. { - std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); + std::lock_guard shared_cache_lock(sf.mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -710,23 +720,32 @@ void fused_attn_arbitrary_seqlen_bwd_impl( using CacheType = std::map; static CacheType sdpa_f16_bprop_cache; - static std::mutex sdpa_f16_bprop_cache_mutex; + static graph_cache::SingleFlight sdpa_f16_bprop_cache_sf; - // Get plan from cache if cache is available, otherwise create one + // Get plan from cache if available; otherwise build it exactly once across + // threads (single-flight), so concurrent misses of the same key don't each + // compile and discard an identical graph. auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - // Lock the map lookup, not the build, so different graphs can build in parallel - graph_and_tensors cached_graph{}; - bool cache_hit = false; + auto &sf = sdpa_f16_bprop_cache_sf; { - std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); + std::unique_lock lock(sf.mutex); + // Wait until the graph is cached, or no other thread is building this key. + sf.cv.wait(lock, [&] { + return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; + }); auto it = cache.find(descriptor); - 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; + if (it != cache.end()) { + graph_and_tensors cached_graph = it->second; // copy under the lock + lock.unlock(); + graph_cache_debug::record_cache_lookup("bwd", /*hit=*/true, cfg, descriptor.device_id); + return cached_graph; + } + // Claim the build for this key, so a concurrent miss waits instead of + // compiling an identical graph. No claim means no waiting. + if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); } + graph_cache_debug::record_cache_lookup("bwd", /*hit=*/false, cfg, descriptor.device_id); + graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -980,10 +999,11 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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. + // Insert our graph. With single-flight we are normally the only builder + // for this key; insert() still tolerates a pre-existing entry and returns + // it. claim_guard releases the build claim and wakes waiters on return. { - std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); + std::lock_guard shared_cache_lock(sf.mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -1343,8 +1363,11 @@ 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.check_forward = true; graph_cfg.derive(); + // Attribute the graph this builds to the support probe, not to a real + // execution: it may be for a config that never runs. + graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { @@ -1368,8 +1391,11 @@ 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.check_forward = false; graph_cfg.derive(); + // Attribute the graph this builds to the support probe, not to a real + // execution: it may be for a config that never runs. + graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 9449f74206..60699cef0a 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -130,23 +130,32 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de // 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; + static graph_cache::SingleFlight sdpa_fp8_fprop_cache_sf; - // Get plan from cache if cache is available, otherwise create one + // Get plan from cache if available; otherwise build it exactly once across + // threads (single-flight), so concurrent misses of the same key don't each + // compile and discard an identical graph. auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - // Lock the map lookup, not the build, so different graphs can build in parallel - graph_and_tensors cached_graph{}; - bool cache_hit = false; + auto& sf = sdpa_fp8_fprop_cache_sf; { - std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); + std::unique_lock lock(sf.mutex); + // Wait until the graph is cached, or no other thread is building this key. + sf.cv.wait(lock, [&] { + return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; + }); auto it = cache.find(descriptor); - 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; + if (it != cache.end()) { + graph_and_tensors cached_graph = it->second; // copy under the lock + lock.unlock(); + graph_cache_debug::record_cache_lookup("fwd", /*hit=*/true, cfg, descriptor.device_id); + return cached_graph; + } + // Claim the build for this key, so a concurrent miss waits instead of + // compiling an identical graph. No claim means no waiting. + if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); } + graph_cache_debug::record_cache_lookup("fwd", /*hit=*/false, cfg, descriptor.device_id); + graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -411,10 +420,11 @@ 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); 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. + // Insert our graph. With single-flight we are normally the only builder + // for this key; insert() still tolerates a pre-existing entry and returns + // it. claim_guard releases the build claim and wakes waiters on return. { - std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); + std::lock_guard shared_cache_lock(sf.mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -626,23 +636,32 @@ void fused_attn_fp8_bwd_impl( using CacheType = std::map; static CacheType sdpa_fp8_bprop_cache; - static std::mutex sdpa_fp8_bprop_cache_mutex; + static graph_cache::SingleFlight sdpa_fp8_bprop_cache_sf; - // Get plan from cache if cache is available, otherwise create one + // Get plan from cache if available; otherwise build it exactly once across + // threads (single-flight), so concurrent misses of the same key don't each + // compile and discard an identical graph. auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - // Lock the map lookup, not the build, so different graphs can build in parallel - graph_and_tensors cached_graph{}; - bool cache_hit = false; + auto& sf = sdpa_fp8_bprop_cache_sf; { - std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); + std::unique_lock lock(sf.mutex); + // Wait until the graph is cached, or no other thread is building this key. + sf.cv.wait(lock, [&] { + return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; + }); auto it = cache.find(descriptor); - 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; + if (it != cache.end()) { + graph_and_tensors cached_graph = it->second; // copy under the lock + lock.unlock(); + graph_cache_debug::record_cache_lookup("bwd", /*hit=*/true, cfg, descriptor.device_id); + return cached_graph; + } + // Claim the build for this key, so a concurrent miss waits instead of + // compiling an identical graph. No claim means no waiting. + if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); } + graph_cache_debug::record_cache_lookup("bwd", /*hit=*/false, cfg, descriptor.device_id); + graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -1037,10 +1056,11 @@ 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); 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. + // Insert our graph. With single-flight we are normally the only builder + // for this key; insert() still tolerates a pre-existing entry and returns + // it. claim_guard releases the build claim and wakes waiters on return. { - std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); + std::lock_guard shared_cache_lock(sf.mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -1385,8 +1405,11 @@ 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.check_forward = true; graph_cfg.derive(); + // Attribute the graph this builds to the support probe, not to a real + // execution: it may be for a config that never runs. + graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { @@ -1411,8 +1434,11 @@ 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.check_forward = false; graph_cfg.derive(); + // Attribute the graph this builds to the support probe, not to a real + // execution: it may be for a config that never runs. + graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 033edf2ead..3961c0871c 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -7,22 +7,42 @@ // ============================================================================ // 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. +// Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG. Two verbosity levels: +// =1 : low volume. Cache event counters, per-build BUILD lines, and the +// end-of-run SUMMARY (aggregate + per-thread) and stage timings. This is +// enough to diagnose redundant rebuilds and profile build cost. +// =2 : high volume (trace). Additionally emits a per-lookup HIT/MISS line with +// the full shorthand config and a per-execution EXEC line. Use only when +// you need to see *which* shapes are hitting/missing -- these fire on +// every cache lookup and execution, so at suite scale they add I/O and +// serialize threads on the stderr lock (perturbing the build timings). +// +// An optional ":" suffix picks which processes emit, defaulting to rank 0 +// so that output does not scale with the world size: "1:all" for every rank, +// "2:0,3" for a specific set. See `rank_selected` for when overriding pays off. // ============================================================================ #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 #include #include +#include +#include +#include +#include +#include + +#include +#include #include "config_and_params.h" @@ -30,23 +50,129 @@ 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 = [] { +// Rank of this process as reported by the launcher, or -1 when there is no +// launcher (a single-process run). First variable that is set wins. +inline int launcher_rank() { + static const int rank = []() -> int { + for (const char *var : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { + const char *v = std::getenv(var); + if (v != nullptr && v[0] != '\0') return std::atoi(v); + } + return -1; + }(); + return rank; +} + +// Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=default, +// 2=trace). Single read at startup, cached. Negligible overhead when unset. +inline int debug_level() { + static const int lvl = [] { const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; + if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; + const int v = std::atoi(e); // stops at the optional ":" suffix + return v > 0 ? v : 1; // any non-empty, non-"0" value enables at least level 1 }(); - return on; + return lvl; } -// More readable, shorter thread IDs (0, 1, 2, ...). +// Whether this process emits diagnostics. Every rank writes to the same stderr, +// so emitting from all of them multiplies the volume by the world size -- and +// under data/tensor parallelism the ranks are running identical shapes, so the +// copies say the same thing. Hence rank 0 only by default. +// +// Context parallelism is the case worth overriding for: the ranks run different +// subsets of the per-step regimes (under p2p, rank 0 never sees the lower-triangle +// config that the last rank does), so their build counts genuinely differ. +// Select with the ":" suffix, e.g. "1:all" or "2:0,3". +inline bool rank_selected() { + static const bool selected = [] { + const int rank = launcher_rank(); + if (rank < 0) return true; // sole process, nothing to filter + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + const char *sep = (e != nullptr) ? std::strchr(e, ':') : nullptr; + if (sep == nullptr) return rank == 0; + const std::string list(sep + 1); + if (list == "all") return true; + for (size_t pos = 0; pos <= list.size();) { + const size_t comma = list.find(',', pos); + const std::string tok = + list.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + if (!tok.empty() && std::atoi(tok.c_str()) == rank) return true; + if (comma == std::string::npos) break; + pos = comma + 1; + } + return false; + }(); + return selected; +} + +// Diagnostics are on at level >= 1, and only for the selected ranks. Unselected +// ranks skip the counters too, so they pay nothing beyond this check. +inline bool enabled() { return debug_level() >= 1 && rank_selected(); } + +// Per-lookup / per-exec trace lines are gated behind level >= 2. +inline bool trace_enabled() { return debug_level() >= 2; } + +// Identifies the emitting process. Distributed PyTorch runs one process per rank +// and they all share this stderr, so without this every line would be ambiguous +// (thread ids restart at 0 in each process). Rank comes from the launcher, if any. +inline const std::string &process_tag() { + static const std::string *tag = [] { + auto *s = new std::string("pid=" + std::to_string(static_cast(::getpid()))); + if (launcher_rank() >= 0) *s += " rank=" + std::to_string(launcher_rank()); + return s; + }(); + return *tag; +} + +// More readable, shorter thread IDs (0, 1, 2, ...). These are assignment order, +// not identity: tid=0 is whichever thread touched this cache first. The one-shot +// THREAD line below maps them to OS thread ids for correlating with nsys/gdb. 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; } +// OS-level thread id, as reported by nsys/gdb/`top -H`. Via syscall rather than +// gettid() so this does not require glibc >= 2.30. +inline int64_t os_thread_id() { return static_cast(::syscall(SYS_gettid)); } + +// True while this thread is inside a backend-support probe (`is_supported_*`), +// which builds a graph speculatively just to answer "is this config supported?". +// Such a build may never be executed -- notably the context-parallel per-step +// probe, which checks regimes this rank never runs -- so it is counted apart +// from builds triggered by a real execution. +inline bool &tl_in_probe() { + static thread_local bool v = false; + return v; +} + +// Marks the calling scope as a support probe. Saves/restores rather than +// clearing, so it stays correct if probes ever nest. +struct ScopedProbe { + bool prev; + ScopedProbe() : prev(tl_in_probe()) { tl_in_probe() = true; } + ~ScopedProbe() { tl_in_probe() = prev; } +}; + +inline const char *src_tag() { return tl_in_probe() ? "probe" : "exec"; } + +// Milliseconds since the first diagnostic event. Used to correlate build +// start/end across threads and detect whether same-shape builds on different +// devices overlap in wall-clock time. +inline double now_ms() { + static const auto t0 = std::chrono::steady_clock::now(); + return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); +} + +// Per-thread build-start timestamp, set at the miss that triggers a build and +// read when the build completes, to report each build's duration. +inline double &tl_build_start_ms() { + static thread_local double v = 0.0; + return v; +} + // Registered at first use. On process exit, prints overall event counters and // graph build timings. inline void register_summary_once(); @@ -63,6 +189,9 @@ inline void register_summary_once(); struct EventCounters { std::atomic built{0}; + // Subset of `built` that a support probe triggered. A large probe share means + // time is going into graphs that may never run. + std::atomic built_probe{0}; std::atomic exec{0}; std::atomic hit{0}; std::atomic miss{0}; @@ -74,27 +203,110 @@ inline EventCounters &counters(bool is_fwd) { 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=%" 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)); +// Per-thread counters, so the summary can break down build/exec/hit/miss by +// thread. In the single-process context-parallel case each device is driven by +// its own thread, so this reveals which thread built/executed what. +struct ThreadCounters { + unsigned tid = 0; + EventCounters fwd; + EventCounters bwd; +}; + +// The registry and its mutex are heap-allocated and deliberately never freed. +// Function-local static destructors and atexit handlers run as a single sequence, +// in reverse order of construction/registration. This registry is built lazily, so +// it can be constructed *after* the summary handler is registered -- in which case +// it would be destroyed *before* that handler runs, leaving the handler to lock a +// destroyed mutex and walk a destroyed vector. Leaking removes the ordering +// question rather than reasoning about it, and the cost is bounded: one mutex and +// one vector for the process, reclaimed by the OS at exit anyway. +inline std::mutex &thread_registry_mutex() { + static std::mutex *m = new std::mutex(); + return *m; +} +inline std::vector &thread_registry() { + static std::vector *v = new std::vector(); + return *v; +} + +// This thread's counter block, leaked for a related but distinct reason: a worker +// thread can exit long before the process does, while the registry keeps a pointer +// to its block for the end-of-run summary. Tying the block's lifetime to the +// thread would leave that pointer dangling. One small struct per thread. +inline ThreadCounters &thread_counters() { + static thread_local ThreadCounters *tc = [] { + auto *p = new ThreadCounters(); + p->tid = thread_seq_id(); + { + std::lock_guard lock(thread_registry_mutex()); + thread_registry().push_back(p); + } + // One line per thread, mapping the short id to something nsys/gdb can match. + std::fprintf(stderr, "[FUSED-ATTN-CACHE] %s | THREAD | tid=%-3u os_tid=%" PRId64 "\n", + process_tag().c_str(), p->tid, os_thread_id()); + std::fflush(stderr); + return p; + }(); + return *tc; +} + +inline EventCounters &thread_counters(bool is_fwd) { + ThreadCounters &tc = thread_counters(); + return is_fwd ? tc.fwd : tc.bwd; +} + +// Format one counter block (aggregate or a single thread's) as one line. +// `tid_field` is the whole thread column, e.g. "tid=3"; the aggregate row passes +// "tid=all" so that it cannot be misread as thread 0's row. +inline std::string format_counter_line(const char *event, const char *tid_field, + const EventCounters &f, const EventCounters &b, + const char *extra) { + char buf[640]; + std::snprintf(buf, sizeof(buf), + "[FUSED-ATTN-CACHE] %s | %-11s | %-7s | fwd miss=%4" PRIu64 ", hit=%4" PRIu64 + ", built=%4" PRIu64 " (for probe %4" PRIu64 "), exec=%4" PRIu64 " | bwd miss=%4" PRIu64 + ", hit=%4" PRIu64 ", built=%4" PRIu64 " (for probe %4" PRIu64 "), exec=%4" PRIu64 "%s\n", + process_tag().c_str(), event, tid_field, f.miss.load(std::memory_order_relaxed), + f.hit.load(std::memory_order_relaxed), f.built.load(std::memory_order_relaxed), + f.built_probe.load(std::memory_order_relaxed), + f.exec.load(std::memory_order_relaxed), b.miss.load(std::memory_order_relaxed), + b.hit.load(std::memory_order_relaxed), b.built.load(std::memory_order_relaxed), + b.built_probe.load(std::memory_order_relaxed), + b.exec.load(std::memory_order_relaxed), extra); + return std::string(buf); +} + +inline void print_counter_block(const char *event, const char *tid_field, const EventCounters &f, + const EventCounters &b, const char *extra = "") { + const std::string line = format_counter_line(event, tid_field, f, b, extra); + std::fputs(line.c_str(), stderr); std::fflush(stderr); } +inline void print_counters(const char *event, const char *extra = "") { + char tid_field[16]; + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); + print_counter_block(event, tid_field, counters(/*is_fwd=*/true), counters(/*is_fwd=*/false), + extra); +} + 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"); + thread_counters(is_fwd).built.fetch_add(1, std::memory_order_relaxed); + if (tl_in_probe()) { + counters(is_fwd).built_probe.fetch_add(1, std::memory_order_relaxed); + thread_counters(is_fwd).built_probe.fetch_add(1, std::memory_order_relaxed); + } + // Report build completion time and this build's wall-clock duration so we can + // tell whether same-shape builds on different devices overlap. + const double t_end = now_ms(); + char extra[80]; + std::snprintf(extra, sizeof(extra), " | src=%-5s t=%.1f dur=%.1f ms", src_tag(), t_end, + t_end - tl_build_start_ms()); + print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD", extra); } inline void record_exec(const char *pass) { @@ -102,17 +314,36 @@ inline void record_exec(const char *pass) { 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"); + thread_counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); + // The per-exec line fires on every execution; keep it out of the level-1 path. + if (!trace_enabled()) return; + char extra[32]; + std::snprintf(extra, sizeof(extra), " | t=%.1f", now_ms()); + print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC", extra); } -inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { +// `device_key` is the cache-scope discriminator from make_cache_key(), not a device ordinal: +// it is the packed (SM arch, SM count) when devices share plans, else the device id. +inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c, + int device_key) { if (!enabled()) return; register_summary_once(); - EventCounters &pc = counters(std::strcmp(pass, "fwd") == 0); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + EventCounters &pc = counters(is_fwd); (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); + EventCounters &tpc = thread_counters(is_fwd); + (hit ? tpc.hit : tpc.miss).fetch_add(1, std::memory_order_relaxed); + const double t = now_ms(); + // A miss triggers a build right after this call; stamp the build start so the + // subsequent BUILD line can report duration. Do this even at level 1. + if (!hit) tl_build_start_ms() = t; + // The per-lookup config dump is the highest-volume line (one per cache probe); + // keep it out of the level-1 path and off the stderr lock unless tracing. + if (!trace_enabled()) return; std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %-3s %-4s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d " + "[FUSED-ATTN-CACHE] %s | %-3s %-4s | tid=%u devkey=%d t=%.1f src=%-5s | 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 @@ -121,9 +352,10 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi " 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), + process_tag().c_str(), pass, hit ? "HIT" : "MISS", thread_seq_id(), device_key, t, src_tag(), + 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.return_max_logit), static_cast(c.check_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), @@ -145,7 +377,6 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi 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); } // ============================================================================ @@ -161,11 +392,14 @@ 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) { +// Bucketed by pass and by whether a support probe drove the build, so the +// summary can say how much of (notably) `build_plans` was speculative. +constexpr size_t kStageBuckets = 4 * static_cast(BuildStage::kCount); +inline StageTiming &stage_timing(bool is_fwd, bool is_probe, BuildStage s) { static std::array table{}; - const size_t idx = - (is_fwd ? 0u : 1u) * static_cast(BuildStage::kCount) + static_cast(s); + const size_t idx = ((is_fwd ? 0u : 1u) * 2u + (is_probe ? 1u : 0u)) * + static_cast(BuildStage::kCount) + + static_cast(s); return table[idx]; } @@ -173,8 +407,10 @@ struct ScopedBuildTimer { BuildStage stage; bool on; bool is_fwd; + bool is_probe; std::chrono::steady_clock::time_point start; - ScopedBuildTimer(bool is_fwd_, BuildStage s) : stage(s), on(enabled()), is_fwd(is_fwd_) { + ScopedBuildTimer(bool is_fwd_, BuildStage s) + : stage(s), on(enabled()), is_fwd(is_fwd_), is_probe(tl_in_probe()) { if (!on) return; register_summary_once(); start = std::chrono::steady_clock::now(); @@ -185,7 +421,7 @@ struct ScopedBuildTimer { static_cast(std::chrono::duration_cast( std::chrono::steady_clock::now() - start) .count()); - StageTiming &t = stage_timing(is_fwd, stage); + StageTiming &t = stage_timing(is_fwd, is_probe, stage); t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); t.calls.fetch_add(1, std::memory_order_relaxed); } @@ -204,23 +440,51 @@ inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; - print_counters("SUMMARY"); + // Build the whole summary in memory and emit it with a single write, so + // that the blocks of concurrently-exiting processes (one per rank under + // torchrun) stay grouped instead of interleaving line by line. + std::string block; + block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary begin =====\n"; + // Per-thread breakdown (sorted by tid). Useful in the single-process + // context-parallel case where each device runs on its own thread. + { + std::lock_guard lock(thread_registry_mutex()); + std::vector blocks = thread_registry(); + std::sort(blocks.begin(), blocks.end(), + [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); + for (const ThreadCounters *tc : blocks) { + char tid_field[16]; + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); + block += format_counter_line("SUMMARY-TID", tid_field, tc->fwd, tc->bwd, ""); + } + } + // Totals last, so they read as the sum of the per-thread lines above. + block += format_counter_line("SUMMARY", "tid=all", counters(/*is_fwd=*/true), + counters(/*is_fwd=*/false), ""); 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=%" PRIu64 - " | time=%9.1f ms | avg=%9.3f ms/call\n", - pass, kStageNames[i], n, total_ms, total_ms / n); + for (int q = 0; q < 2; ++q) { + const bool is_probe = (q == 1); + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + const BuildStage s = static_cast(i); + const StageTiming &t = stage_timing(is_fwd, is_probe, 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; + char line[288]; + std::snprintf(line, sizeof(line), + "[FUSED-ATTN-CACHE] %s | %-3s src=%-5s %-22s | calls=%" PRIu64 + " | time=%9.3f ms/call\n", + process_tag().c_str(), pass, is_probe ? "probe" : "exec", kStageNames[i], + n, total_ms / n); + block += line; + } } } + block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary end =====\n"; + std::fwrite(block.data(), 1, block.size(), stderr); std::fflush(stderr); }); return true; @@ -229,6 +493,74 @@ inline void register_summary_once() { } } // namespace graph_cache_debug + +// ============================================================================ +// Single-flight graph cache coordination. +// +// The fused-attention graph caches are process-wide and shared across threads. +// The lock is intentionally released while a graph is compiled so that +// *different* graphs can build in parallel. The downside is that when several +// threads miss the *same* key at the same instant (e.g. the device-worker +// threads of a single-process context-parallel run stepping in lockstep), they +// all compile an identical graph and all but one discard the result at insert +// time -- wasted host-side `build_plans` work. +// +// A single-flight (a.k.a. "thundering herd") guard closes that gap: at most one +// thread compiles a given key while the others wait for it. Distinct keys still +// build concurrently, so the parallel-build win is kept. +// +// Usage in a get_graph path: +// static SingleFlight sf; +// { +// std::unique_lock lock(sf.mutex); +// sf.cv.wait(lock, [&] { +// return cache.count(key) != 0 || sf.in_progress.count(key) == 0; +// }); +// if (auto it = cache.find(key); it != cache.end()) { ...cache hit... } +// sf.in_progress.insert(key); // claim the build +// } +// ClaimGuard guard{sf, key}; // auto-release + notify +// ...build... +// { std::lock_guard lock(sf.mutex); cache.insert({key, graph}); } +// ============================================================================ +namespace graph_cache { + +// Opt-in with NVTE_FUSED_ATTN_CACHE_SINGLE_FLIGHT=1. When off, no thread claims a +// key, so the wait below falls through immediately and concurrent misses of one +// key each compile, with all but one result discarded at insert: wasted host work, +// but no thread ever blocks on another thread's compile. +inline bool single_flight_enabled() { + static const bool on = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_SINGLE_FLIGHT"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return on; +} + +template +struct SingleFlight { + std::mutex mutex; + std::condition_variable cv; + std::set in_progress; +}; + +// RAII: on scope exit, drop this thread's build claim on `key` and wake any +// threads waiting on the same key. Clears the claim even if the build throws, +// so waiters never deadlock (they simply re-elect a builder). +template +struct ClaimGuard { + SingleFlight &sf; + const KeyT &key; + ~ClaimGuard() { + { + std::lock_guard lock(sf.mutex); + sf.in_progress.erase(key); + } + sf.cv.notify_all(); + } +}; + +} // namespace graph_cache } // namespace fused_attn } // 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 938fa1747e..0369ed62bf 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -413,6 +413,17 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, const void *buf, size_t size_in_bytes); +/*! \brief Prefix on the ``nvte_get_fused_attn_backend_v2`` diagnostic when a configuration is + * rejected solely by the backward-pass support check. + * + * Its presence tells the caller that the forward pass is supported and only the backward pass + * is not, so re-querying the same configuration with ``is_training = false`` may succeed. Its + * absence means the forward pass itself was rejected, for which dropping to inference cannot + * help. Mirrored on the Python side as + * ``transformer_engine.pytorch.attention.dot_product_attention.utils.FUSED_ATTN_BWD_REJECT_PREFIX``. + */ +#define NVTE_FUSED_ATTN_BWD_REJECT_PREFIX "[backward] " + /*! \brief Get fused-attention backend based on user configuration. * * This function passes the user configuration to cuDNN frontend, runs its support checks, @@ -422,7 +433,9 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, * \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 explaining why there is no support. + * if not, a diagnostic message explaining why there is no support, + * prefixed with ``NVTE_FUSED_ATTN_BWD_REJECT_PREFIX`` when only the + * backward pass is unsupported. * 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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/__init__.py b/transformer_engine/pytorch/attention/dot_product_attention/__init__.py index 941f94f105..bd54ca3134 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/__init__.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/__init__.py @@ -4,6 +4,18 @@ """Python interface for dot product attention""" -from .dot_product_attention import DotProductAttention, _attention_backends +from .dot_product_attention import ( + DotProductAttention, + BackendSelectionProbe, + DryRunResult, + dry_run_backend_selection, + _attention_backends, +) -__all__ = ["DotProductAttention", "_attention_backends"] +__all__ = [ + "DotProductAttention", + "BackendSelectionProbe", + "DryRunResult", + "dry_run_backend_selection", + "_attention_backends", +] 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 93296e62dd..cd9555da52 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4946,7 +4946,8 @@ def cp_per_step_configs( 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, t_q, t_kv): + def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv, window=None): + w_left, w_right = window if window is not None else (window_left, window_right) return { "attn_mask_type": mask, "max_seqlen_q": s_q, @@ -4955,8 +4956,8 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): "num_tokens_kv": t_kv, "num_attn_heads": heads, "num_gqa_groups": gqa, - "window_size_left": window_left, - "window_size_right": window_right, + "window_size_left": w_left, + "window_size_right": w_right, "bottom_right_diagonal": bottom_right, } @@ -4982,20 +4983,16 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): 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 - t_q = num_tokens_q // 2 + # Each step narrows max_seqlen_*, but the token counts it dispatches with are the + # rank's full Q tokens and the all-gathered KV tokens, unchanged across steps. + # Scaling them per step would key the probe's graph differently from the one the + # step looks up, and rebuild every graph this probes at execution time. + t_q = num_tokens_q + t_kv = num_tokens_kv * cp_size # 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, - t_q, - num_tokens_kv * cp_size * s_kv // max_seqlen_kv if max_seqlen_kv else 0, - ) + config(mask, s_q, s_kv, num_heads, num_gqa_groups, br, t_q, t_kv) for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) ] 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 3c079cc5d5..2117c528c7 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 @@ -3,9 +3,11 @@ # See LICENSE for license information. """Attention.""" -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext +import dataclasses import math import os +import threading from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings import logging @@ -72,9 +74,122 @@ "use_fused_attention": None, "fused_attention_backend": None, "use_unfused_attention": None, + "available_backends": None, + "fused_attn_reject_reason": None, "backend_selection_requires_update": False, } +# Dry-run backend selection: see dry_run_backend_selection(). +_dpa_dry_run = threading.local() + + +class _DryRunComplete(Exception): + """Unwinds the forward pass once backend selection is known.""" + + +@dataclasses.dataclass +class BackendSelectionProbe: + """Backend selection at a single `DotProductAttention` site. + + Attributes + ---------- + attention_params : AttentionParams + The parameters the module actually resolved, i.e. the exact cache key the real run + would use. + available_backends : List[bool] + [flash, fused, unfused] support for this configuration. + fused_attention_backend : Optional[FusedAttnBackend] + The selected `FusedAttention` sub-backend, or `None`. + fused_attn_reject_reason : Optional[str] + Why `FusedAttention` was ruled out, or `None`. A reason starting with + `FUSED_ATTN_BWD_REJECT_PREFIX` means only the backward pass is unsupported. + """ + + attention_params: Optional["dpa_utils.AttentionParams"] = None + available_backends: Optional[List[bool]] = None + fused_attention_backend: Optional[Any] = None + fused_attn_reject_reason: Optional[str] = None + + +@dataclasses.dataclass +class DryRunResult: + """Backend selection across every `DotProductAttention` site a dry run reached. + + A module runs only if all of its attention sites are supported, so the `*_supported` + properties require every recorded site to support the backend. `probes` holds the + per-site detail, in the order the sites were reached. + """ + + stop_after: int = 1 + probes: List[BackendSelectionProbe] = dataclasses.field(default_factory=list) + + def _all_support(self, index: int) -> bool: + return bool(self.probes) and all(p.available_backends[index] for p in self.probes) + + @property + def flash_supported(self) -> bool: + """Whether `FlashAttention` supports every attention site reached.""" + return self._all_support(0) + + @property + def fused_supported(self) -> bool: + """Whether `FusedAttention` supports every attention site reached.""" + return self._all_support(1) + + @property + def unfused_supported(self) -> bool: + """Whether `UnfusedDotProductAttention` supports every attention site reached.""" + return self._all_support(2) + + @property + def fused_attn_reject_reason(self) -> Optional[str]: + """Why `FusedAttention` was ruled out, from the first site that ruled it out.""" + return next( + (p.fused_attn_reject_reason for p in self.probes if p.fused_attn_reject_reason), + None, + ) + + +@contextmanager +def dry_run_backend_selection(stop_after: int = 1): + """Resolve attention backends for a module without running attention. + + Call a module as usual inside this context. `DotProductAttention.forward` resolves its + configuration exactly as a real run would, records which backends support it, and then + unwinds the forward pass before any attention is executed. This avoids having to + predict the configuration a module will produce, e.g. the `qkv_layout` that + `MultiheadAttention` derives from its packed projection output. + + Because the resolved configuration is identical to the real run's, the backend query + populates the same cache entries the real run will hit. + + Parameters + ---------- + stop_after : int, default = 1 + Unwind after this many attention sites have been recorded. Modules with several + attention sites, e.g. a `TransformerLayer` with `layer_type="decoder"`, need a + higher value to reach the later ones. Note that the sites before the last one + execute for real, since that is the only way to arrive at what follows them. + + .. code-block:: python + + with dry_run_backend_selection() as dry_run: + model(hidden_states, attn_mask_type="causal") + if not dry_run.fused_supported: + print(dry_run.fused_attn_reject_reason) + """ + assert stop_after >= 1, "stop_after must be at least 1" + result = DryRunResult(stop_after=stop_after) + previous = getattr(_dpa_dry_run, "result", None) + _dpa_dry_run.result = result + try: + yield result + except _DryRunComplete: + pass + finally: + _dpa_dry_run.result = previous + + _alibi_cache = { "_num_heads": None, "_alibi_slopes": None, @@ -2041,6 +2156,9 @@ def forward( use_flash_attention = False use_fused_attention = False use_unfused_attention = True + fused_attention_backend = None + available_backends = [False, False, True] + fused_attn_reject_reason = None else: if ( _attention_backends["attention_params"] is None @@ -2055,7 +2173,8 @@ def forward( use_fused_attention, fused_attention_backend, use_unfused_attention, - _, + available_backends, + fused_attn_reject_reason, ) = dpa_utils.get_attention_backend(attention_params) # Set global _attention_backends var using return value # from get_attention_backend() @@ -2064,6 +2183,8 @@ def forward( _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["available_backends"] = available_backends + _attention_backends["fused_attn_reject_reason"] = fused_attn_reject_reason _attention_backends["backend_selection_requires_update"] = False if use_flash_attention: self.logger.info( @@ -2083,6 +2204,21 @@ def forward( use_fused_attention = _attention_backends["use_fused_attention"] fused_attention_backend = _attention_backends["fused_attention_backend"] use_unfused_attention = _attention_backends["use_unfused_attention"] + available_backends = _attention_backends["available_backends"] + fused_attn_reject_reason = _attention_backends["fused_attn_reject_reason"] + + dry_run = getattr(_dpa_dry_run, "result", None) + if dry_run is not None: + dry_run.probes.append( + BackendSelectionProbe( + attention_params=attention_params, + available_backends=available_backends, + fused_attention_backend=fused_attention_backend, + fused_attn_reject_reason=fused_attn_reject_reason, + ) + ) + if len(dry_run.probes) >= dry_run.stop_after: + raise _DryRunComplete # 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 eb10f0c03d..f37bdacefc 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -73,6 +73,11 @@ _cu_seqlens_cache = {} +# Mirrors NVTE_FUSED_ATTN_BWD_REJECT_PREFIX in common/include/transformer_engine/fused_attn.h. +# A rejection reason carrying this prefix means only the backward pass is unsupported, so the +# same config may be supported with is_training=False. Keep the two definitions in sync. +FUSED_ATTN_BWD_REJECT_PREFIX = "[backward] " + class AttentionLogging: """ @@ -466,6 +471,11 @@ def get_attention_backend( available_backends : List[bool] 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]. + fused_attention_reject_reason : Optional[str] + Why `FusedAttention` was ruled out, or `None` if it was selected or was never queried + (e.g. disabled via `NVTE_FUSED_ATTN=0`). A reason starting with + `FUSED_ATTN_BWD_REJECT_PREFIX` means only the backward pass is unsupported, so the same + config may be supported with `is_training=False`. """ # 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 @@ -1531,6 +1541,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # Filter: cuDNN support fused_attention_backend = None + fused_attention_reject_reason = None if use_fused_attention: 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" @@ -1642,10 +1653,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt 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 + if fu_core_attention_bias_shape is not None: + 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 # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. @@ -1662,6 +1674,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_fused_attention = False fused_attention_backend = None + fused_attention_reject_reason = reject_message break if ( @@ -1722,12 +1735,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): - logger.debug( - "Disabling FusedAttention for determinism reasons with FP8 on arch < sm90 or cuDNN" - " < 9.19.0" - ) + reason = "determinism with FP8 is not supported on arch < sm90 or cuDNN < 9.19.0" + logger.debug("Disabling FusedAttention for %s", reason) use_fused_attention = False fused_attention_backend = None + fused_attention_reject_reason = FUSED_ATTN_BWD_REJECT_PREFIX + reason if ( fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] and is_training @@ -1737,9 +1749,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt or cudnn_version < (8, 9, 5) ) ): - logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") + reason = "determinism with post_scale_bias is not supported" + logger.debug("Disabling FusedAttention for %s", reason) use_fused_attention = False fused_attention_backend = None + fused_attention_reject_reason = FUSED_ATTN_BWD_REJECT_PREFIX + reason # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 @@ -1849,6 +1863,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fused_attention_backend, use_unfused_attention, available_backends, + fused_attention_reject_reason, ) From e027294c552dc87521afea66b74e090c99269a1b Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:22:27 -0700 Subject: [PATCH 71/88] Revert "temporary changes: cache debug, timers, single flight, is_probe, dry-run, still build plans in probes" This reverts commit 8fdd81d291a96ef9aebc0136894b90cb040eed0c. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 236 +++++----- tests/pytorch/test_torch_compile.py | 1 - tests/pytorch/utils.py | 63 +-- .../common/fused_attn/config_and_params.cpp | 89 +--- .../common/fused_attn/config_and_params.h | 17 +- .../common/fused_attn/fused_attn.cpp | 28 +- .../fused_attn_f16_arbitrary_seqlen.cu | 90 ++-- .../common/fused_attn/fused_attn_fp8.cu | 90 ++-- .../common/fused_attn/graph_cache_debug.h | 428 ++---------------- .../include/transformer_engine/fused_attn.h | 15 +- .../dot_product_attention/__init__.py | 16 +- .../dot_product_attention/context_parallel.py | 25 +- .../dot_product_attention.py | 140 +----- .../attention/dot_product_attention/utils.py | 33 +- 14 files changed, 273 insertions(+), 998 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2b9b64a026..ca93fecc9e 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -29,7 +29,6 @@ _attention_backends, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( - FUSED_ATTN_BWD_REJECT_PREFIX, FlashAttentionUtils, check_set_window_size, ) @@ -63,7 +62,6 @@ ModelConfig, dtype_tols, get_available_attention_backends, - probe_attention_backends, ) # Check if hardware supports FP8 attention. @@ -105,18 +103,18 @@ def reset_global_fp8_state(): model_configs_base = { # test: ModelConfig(b, sq, hq, dqk) - "base_1_0": ModelConfig(8, 131072, 16, 64), + "base_1_0": ModelConfig(8, 128, 16, 64), "base_1_1": ModelConfig(4, 128, 16, 64, max_seqlen_kv=256), "base_2_0": ModelConfig(2, 2048, 24, 128), "base_2_1": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096), - # "base_3_0": ModelConfig(8, 1, 16, 128, max_seqlen_kv=2048), - # "base_3_1": ModelConfig(8, 1, 16, 256, max_seqlen_kv=2048), - # "base_4_0": ModelConfig(8, 1, 16, 192, max_seqlen_kv=2048), - # "base_4_1": ModelConfig(8, 128, 16, 192, max_seqlen_kv=2048), - # "base_5_0": ModelConfig(8, 1, 16, 512, max_seqlen_kv=2048), - # "base_5_1": ModelConfig(8, 128, 16, 512, max_seqlen_kv=2048), - # "base_6_0": ModelConfig(8, 1, 16, 1024, max_seqlen_kv=2048), - # "base_6_1": ModelConfig(8, 128, 16, 1024, max_seqlen_kv=2048), + "base_3_0": ModelConfig(8, 1, 16, 128, max_seqlen_kv=2048), + "base_3_1": ModelConfig(8, 1, 16, 256, max_seqlen_kv=2048), + "base_4_0": ModelConfig(8, 1, 16, 192, max_seqlen_kv=2048), + "base_4_1": ModelConfig(8, 128, 16, 192, max_seqlen_kv=2048), + "base_5_0": ModelConfig(8, 1, 16, 512, max_seqlen_kv=2048), + "base_5_1": ModelConfig(8, 128, 16, 512, max_seqlen_kv=2048), + "base_6_0": ModelConfig(8, 1, 16, 1024, max_seqlen_kv=2048), + "base_6_1": ModelConfig(8, 128, 16, 1024, max_seqlen_kv=2048), } @@ -180,20 +178,17 @@ def test_dot_product_attention( "Setting is_training to False as cuDNN does not support dbias for" f" {config.bias_shape=} " ) - available_backends, _, fused_attn_backends, reject_reason = get_available_attention_backends( + available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtype, qkv_layout=qkv_layout, pad_between_seqs=pad_between_seqs, is_training=is_training, deterministic=_deterministic, - return_reason=True, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - # Retry in inference mode only when the backward pass alone is what fused attention does not - # support; for any other reason dropping is_training cannot make it available. - if not fused_attn_supported and (reject_reason or "").startswith(FUSED_ATTN_BWD_REJECT_PREFIX): + if not fused_attn_supported: is_training = False available_backends, _, fused_attn_backends = get_available_attention_backends( config, @@ -210,7 +205,6 @@ def test_dot_product_attention( pytest.skip("Less than two backends to compare.") # UnfusedDotProductAttention backend - unfused_attn_supported=False if unfused_attn_supported: unfused_attn_fwd, unfused_max_logit, unfused_attn_bwd = _run_dot_product_attention( dtype, @@ -236,7 +230,6 @@ def test_dot_product_attention( ) # FlashAttention backend - flash_attn_supported = False if flash_attn_supported: flash_attn_fwd, _, flash_attn_bwd = _run_dot_product_attention( dtype, @@ -1591,40 +1584,35 @@ def test_transformer_layer( config = model_configs[model] tols = dict(atol=5e-2, rtol=5e-2) - # Test backend availability. Dry-run the module under test so the query uses the exact - # configuration it resolves, rather than a restatement of it that can drift. A decoder - # layer runs self-attention and then cross-attention, and needs both to be supported. + # Test backend availability is_training = True - num_attn_sites = 2 if config.attn_type == "cross" else 1 - - def probe(is_training): - return probe_attention_backends( - _run_transformer_layer, - dtype, + available_backends, _, fused_attn_backends = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout=( + qkv_format.replace("hd", "h3d") if fused_qkv_params else qkv_format.replace("hd", "3hd") + ), + is_training=is_training, + deterministic=_deterministic, + ) + flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends + if not fused_attn_supported: + is_training = False + available_backends, _, fused_attn_backends = get_available_attention_backends( config, - "", - ckpt_attn, - qkv_format, - fused_qkv_params, - RoPE, - is_training, - num_attn_sites=num_attn_sites, + qkv_dtype=dtype, + qkv_layout=( + qkv_format.replace("hd", "h3d") + if fused_qkv_params + else qkv_format.replace("hd", "3hd") + ), + is_training=is_training, + deterministic=_deterministic, ) - - dry_run = probe(is_training) - # Retry in inference mode only when the backward pass alone is unsupported (see - # test_dot_product_attention). - if not dry_run.fused_supported and (dry_run.fused_attn_reject_reason or "").startswith( - FUSED_ATTN_BWD_REJECT_PREFIX - ): - is_training = False - dry_run = probe(is_training) - flash_attn_supported = dry_run.flash_supported - fused_attn_supported = dry_run.fused_supported - unfused_attn_supported = dry_run.unfused_supported + flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends # Skip if only unfused backend is supported - if (fused_attn_supported + flash_attn_supported + unfused_attn_supported) < 2: + if (len(fused_attn_backends) + flash_attn_supported + unfused_attn_supported) < 2: pytest.skip("Less than two backends to compare.") # Skip if qkv_format = thd and "padding" not in attn_mask_type if qkv_format == "thd" and "padding" not in config.attn_mask_type: @@ -1657,7 +1645,6 @@ def probe(is_training): ) # FlashAttention backend - flash_attn_supported = False if flash_attn_supported: flash_attn_fwd, flash_attn_bwd = _run_transformer_layer( dtype, @@ -1740,13 +1727,17 @@ def _run_transformer_layer( ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Run TransformerLayer module with one forward pass and one backward pass""" - # Set RNG and environment variables. An empty `backend` leaves the caller's choice of - # enabled backends alone, so probe_attention_backends() can ask about all of them. + # Set RNG and environment variables reset_rng_states() - if backend: - os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FlashAttention" else "0" - os.environ["NVTE_FUSED_ATTN"] = "1" if backend == "FusedAttention" else "0" - os.environ["NVTE_UNFUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention" else "0" + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + if backend == "FlashAttention": + os.environ["NVTE_FLASH_ATTN"] = "1" + if backend == "FusedAttention": + 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 @@ -2082,50 +2073,49 @@ def get_model(dtype, config): attn_mask_type = "causal" model_configs_fp8_vs_f16 = { # test: ModelConfig(b, sq, hq, dqk) - # "fp8_9": ModelConfig( - # 2, - # 2048, - # 128, - # 192, - # head_dim_v=128, - # ), - # "fp8_10": ModelConfig( - # 2, - # 2048, - # 128, - # 192, - # head_dim_v=128, - # attn_mask_type="causal", - # ), - # "fp8_11": ModelConfig( - # 2, - # 2048, - # 128, - # 192, - # head_dim_v=128, - # attn_mask_type="causal_bottom_right", - # ), - # "fp8_12": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), - # "fp8_13": ModelConfig( - # 2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(128, 0) - # ), - # "fp8_14": ModelConfig(2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), - "fp8_15": ModelConfig(1, 8192, 64, 64, #attn_mask_type="causal", #window_size=(128, 0) - ), - # "fp8_16": ModelConfig( - # 1, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" - # ), - # "fp8_17": ModelConfig( - # 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" - # ), - # "fp8_18": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), - # "fp8_19": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), - # "fp8_20": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), + "fp8_9": ModelConfig( + 2, + 2048, + 128, + 192, + head_dim_v=128, + ), + "fp8_10": ModelConfig( + 2, + 2048, + 128, + 192, + head_dim_v=128, + attn_mask_type="causal", + ), + "fp8_11": ModelConfig( + 2, + 2048, + 128, + 192, + head_dim_v=128, + attn_mask_type="causal_bottom_right", + ), + "fp8_12": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), + "fp8_13": ModelConfig( + 2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(128, 0) + ), + "fp8_14": ModelConfig(2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), + "fp8_15": ModelConfig(1, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0)), + "fp8_16": ModelConfig( + 1, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" + ), + "fp8_17": ModelConfig( + 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" + ), + "fp8_18": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), + "fp8_19": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), + "fp8_20": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), } -param_types_fp8_vs_f16 = [torch.bfloat16] #[torch.float16, torch.bfloat16] -qkv_layout_fp8_vs_f16 = ["sbhd_sbhd_sbhd"] #["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"] -qkv_format_fp8_vs_f16 = ["sbhd"] #["bshd", "sbhd"] +param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16] +qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"] +qkv_format_fp8_vs_f16 = ["bshd", "sbhd"] @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") @@ -2133,11 +2123,11 @@ def get_model(dtype, config): @pytest.mark.parametrize("dtype", param_types_fp8_vs_f16) @pytest.mark.parametrize("model", model_configs_fp8_vs_f16.keys()) @pytest.mark.parametrize("qkv_format", qkv_format_fp8_vs_f16) -@pytest.mark.parametrize("input_layernorm", [False]) #True, False]) -@pytest.mark.parametrize("fp8_dpa_bwd", [True])#, False]) -@pytest.mark.parametrize("RoPE", [False]) #True, False]) -@pytest.mark.parametrize("is_training", [True]) #, False]) -@pytest.mark.parametrize("scaling_mode", ["delayed"]) #, "current", "mxfp8"]) +@pytest.mark.parametrize("input_layernorm", [True, False]) +@pytest.mark.parametrize("fp8_dpa_bwd", [True, False]) +@pytest.mark.parametrize("RoPE", [True, False]) +@pytest.mark.parametrize("is_training", [True, False]) +@pytest.mark.parametrize("scaling_mode", ["delayed", "current", "mxfp8"]) def test_mha_fp8_vs_f16( dtype, model, @@ -2179,37 +2169,29 @@ def test_mha_fp8_vs_f16( ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe - # Dry-run the module under test so the query uses the exact configuration it resolves, - # rather than a restatement of it that can drift. - fp8_probe = probe_attention_backends( - _run_mha_fp8_vs_f16, - dtype, + available_backends, _, _ = get_available_attention_backends( config, - True, - qkv_format, - input_layernorm, - RoPE, - is_training, - fp8_recipe, + qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, + qkv_layout=qkv_format.replace("hd", "h3d"), + fp8=True, + fp8_meta=fp8_meta, + is_training=is_training, + deterministic=_deterministic, ) - flash_attn_supported = fp8_probe.flash_supported - fused_attn_supported_fp8 = fp8_probe.fused_supported - f16_probe = probe_attention_backends( - _run_mha_fp8_vs_f16, - dtype, + flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends + available_backends, _, fused_attn_backends = get_available_attention_backends( config, - False, - qkv_format, - input_layernorm, - RoPE, - is_training, - fp8_recipe, + qkv_dtype=dtype, + qkv_layout=qkv_format.replace("hd", "h3d"), + is_training=is_training, + deterministic=_deterministic, ) - fused_attn_supported_f16 = f16_probe.fused_supported + _, fused_attn_supported_f16, _ = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: - pytest.skip(fp8_probe.fused_attn_reject_reason or "No FP8 attention backend available.") + pytest.skip("No FP8 attention backend available.") if not fused_attn_supported_f16: - pytest.skip(f16_probe.fused_attn_reject_reason or "No reference backend available.") + pytest.skip("No reference backend available.") if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" @@ -2318,8 +2300,6 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attention_type="self", qkv_weight_interleaved=True, qkv_format=qkv_format, - window_size=config.window_size, - softmax_type=config.softmax_type, ).to(dtype=dtype, device="cuda") if not is_training: mha = mha.eval() diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 05f8a36c4a..dc0da5106e 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -583,7 +583,6 @@ def fn(x, params): fused_attention_backend, use_unfused_attention, _, - _, ) = dpa_utils.get_attention_backend(params) # Encode the full selection (enabled backends + fused sub-backend) in # the tensor value: without a tensor op dynamo skips the frame entirely diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 4652d433e7..90fbcc16b5 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -20,11 +20,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch import InferenceParams, QuantizedTensor from transformer_engine.pytorch import DType -from transformer_engine.pytorch.attention.dot_product_attention import ( - DryRunResult, - dry_run_backend_selection, - _attention_backends, -) +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, AttentionParams, @@ -340,45 +336,6 @@ def logging_context(highest_level=logging.WARNING): logging.disable(previous_level) -def probe_attention_backends(run_fn, *args, num_attn_sites: int = 1, **kwargs) -> DryRunResult: - """Which backends support the configuration a module actually produces. - - `run_fn(*args, **kwargs)` should build and call the module under test exactly as a real - run would. It is aborted inside `DotProductAttention`, once backend selection is known - but before any attention executes. - - Prefer this over `get_available_attention_backends()` for module-level tests. The - latter needs the caller to restate the configuration, which means predicting what the - module derives internally, e.g. the `qkv_layout` that `MultiheadAttention` gets from - its packed projection output. Whenever such a prediction drifts from the module, cuDNN - builds a graph under a cache key that is never executed. - - All three backends are enabled before `run_fn` is called, so a run function that does - not force a backend itself is probed against all of them. A run function that does - force one, e.g. `_run_transformer_layer(backend=...)`, is instead probed under exactly - the environment it will really use; read the matching `*_supported` property. - - Set `num_attn_sites` above 1 for a module that reaches `DotProductAttention` more than - once, e.g. a `TransformerLayer` with `layer_type="decoder"`, which runs self-attention - and then cross-attention. The `*_supported` properties then require every site to be - supported, which is what the module needs to run. - """ - os.environ["NVTE_FLASH_ATTN"] = "1" - os.environ["NVTE_FUSED_ATTN"] = "1" - os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - - with dry_run_backend_selection(stop_after=num_attn_sites) as dry_run: - run_fn(*args, **kwargs) - - assert len(dry_run.probes) == num_attn_sites, ( - f"dry run reached {len(dry_run.probes)} attention site(s), expected" - f" {num_attn_sites}; check num_attn_sites" - ) - _attention_backends["backend_selection_requires_update"] = True - return dry_run - - def get_available_attention_backends( config: ModelConfig, qkv_dtype: torch.dtype, @@ -395,7 +352,6 @@ def get_available_attention_backends( cp_size: int = 1, cp_size_a2a: int = 1, skip_fused_attn: bool = False, - return_reason: bool = False, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration @@ -403,11 +359,6 @@ def get_available_attention_backends( fused-attention backends are then empty, while the FlashAttention and unfused results are unaffected. This skips cuDNN's support checks, which build and cache a graph per configuration. - - Set `return_reason=True` to append the fused-attention rejection reason to the returned - tuple. A reason starting with `FUSED_ATTN_BWD_REJECT_PREFIX` means only the backward pass - is unsupported, i.e. re-querying with `is_training=False` may report fused attention as - available; any other reason means it will not. """ os.environ["NVTE_FLASH_ATTN"] = "1" @@ -491,7 +442,6 @@ def test(): fused_attention_backend, use_unfused_attention, available_backends, - fused_attention_reject_reason, ) = get_attention_backend(attention_params) # Check if FA3 is an available backend when num_splits != 1 if available_backends[0]: @@ -505,23 +455,16 @@ def test(): _attention_backends["fused_attention_backend"] = fused_attention_backend _attention_backends["use_unfused_attention"] = use_unfused_attention _attention_backends["backend_selection_requires_update"] = False - return ( - available_backends, - flash_attention_backend, - fused_attention_backend, - fused_attention_reject_reason, - ) + return available_backends, flash_attention_backend, fused_attention_backend backends = {1: "F16_arbitrary_seqlen", 2: "FP8"} if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() _attention_backends["backend_selection_requires_update"] = True - available_backends, flash_attention_backend, fused_attention_backend, reject_reason = test() + 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) - if return_reason: - return available_backends, flash_attention_backend, fused_attn_backends, reject_reason return available_backends, flash_attention_backend, 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 358beb7641..ca4214dac3 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -9,13 +9,10 @@ #include #include -#include -#include #include #include "../common.h" #include "../util/cuda_runtime.h" -#include "graph_cache_debug.h" namespace { @@ -27,69 +24,6 @@ void uint8_to_bool(const void *in, bool &out) { out = static_cast(*reinterpret_cast(in)); } -// Whether all visible CUDA devices can execute each other's cuDNN graphs. A plan is compiled -// against both the SM architecture and the SM count, and cuDNN requires both to match for a -// graph built on one device to run on another, so a difference in either one rules out sharing. -// SM counts differ across devices in practice even at a fixed arch, e.g. when MIG partitions -// or a harvested SKU are mixed in. -bool all_visible_devices_share_plans() { - const int n = transformer_engine::cuda::num_devices(); - if (n <= 1) return true; - const int arch0 = transformer_engine::cuda::sm_arch(0); - const int sm_count0 = transformer_engine::cuda::sm_count(0); - for (int i = 1; i < n; ++i) { - if (transformer_engine::cuda::sm_arch(i) != arch0) return false; - if (transformer_engine::cuda::sm_count(i) != sm_count0) return false; - } - return true; -} - -// Width reserved for the SM count in the packed cache key below. -constexpr int kSmCountBits = 16; - -// Scope of the fused-attention graph cache across devices in a single process. -// >= 0 : shared key packing (SM arch, SM count) -- all devices reuse one graph per shape -// (homogeneous node; a cuDNN plan compiled for this arch and SM count is valid on -// every device). -// -1 : per-device -- key by device id (heterogeneous node, or forced off). -// Computed once. The two schemes are never mixed within a process, so the packed values need -// not avoid the small device ids. -// -// NVTE_FUSED_ATTN_CACHE_PER_DEVICE=1 forces per-device keying. It is a debug-only -// escape hatch, not a tuning knob: the homogeneity check above is what keeps a -// mixed-arch node correct, so the only reasons to set it are A/B comparison -// against the old behavior, or working around a cuDNN plan-portability bug in the -// field without a rebuild. -int fused_attn_cache_arch_key() { - static const int key = [] () -> int { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_PER_DEVICE"); - const bool force_per_device = (e != nullptr && e[0] != '\0' && e[0] != '0'); - const int n = transformer_engine::cuda::num_devices(); - const bool per_device = force_per_device || !all_visible_devices_share_plans(); - const int arch = per_device ? -1 : transformer_engine::cuda::sm_arch(0); - const int sm_count = per_device ? -1 : transformer_engine::cuda::sm_count(0); - const int result = per_device ? -1 : ((arch << kSmCountBits) | sm_count); - // One-shot: make the resolved cache scope unambiguous in the diagnostics. - if (transformer_engine::fused_attn::graph_cache_debug::enabled()) { - const char *tag = transformer_engine::fused_attn::graph_cache_debug::process_tag().c_str(); - if (result >= 0) { - std::fprintf(stderr, - "\n[FUSED-ATTN-CACHE] %s | cache scope = arch %d + %d SM(s) (key %d) shared " - "across %d device(s)\n", - tag, arch, sm_count, result, n); - } else { - std::fprintf(stderr, - "\n[FUSED-ATTN-CACHE] %s | cache scope = per-device across %d device(s) (%s)\n", - tag, n, force_per_device ? "forced by NVTE_FUSED_ATTN_CACHE_PER_DEVICE" - : "devices differ in arch or SM count"); - } - std::fflush(stderr); - } - return result; - }(); - return key; -} - } // namespace namespace transformer_engine { @@ -162,14 +96,8 @@ void FusedAttnConfig::derive() { FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig cache_cfg = *this; - // Scope the graph cache across devices in a single process. cuDNN plans are compiled against - // a specific SM architecture and SM count, so on a node whose devices agree on both the plan - // built on one device is valid on all of them: key by (arch, SM count) so a shape is built - // once and shared (single-flight collapses the otherwise-per-device duplicate builds). - // When the devices differ in either one (or keying is forced off) fall back to per-device - // keying, since a plan is not portable in that case. - const int arch_key = fused_attn_cache_arch_key(); - cache_cfg.device_id = (arch_key >= 0) ? arch_key : cuda::current_device(); + // 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; @@ -193,7 +121,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { } cache_cfg.num_tokens_q = 0; cache_cfg.num_tokens_kv = 0; - const bool bucket_batch = !check_forward || !cache_cfg.uses_cu_seqlens_directly; + const bool bucket_batch = !is_forward || !cache_cfg.uses_cu_seqlens_directly; if (bucket_batch) { cache_cfg.batch_size = cache_cfg.bucketed_batch_size; } @@ -205,7 +133,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { // Restrict each direction's key to the fields its graph actually consumes, so // no redundant graphs are built and no cache misses either - if (check_forward) { + if (is_forward) { cache_cfg.do_dtype = kNVTEBFloat16; cache_cfg.dqkv_dtype = kNVTEBFloat16; cache_cfg.do_format = NVTE_QKV_Format_NOT_SET; @@ -222,10 +150,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig FusedAttnFwdParams::make_config() const { const FusedAttnFwdParams ¶ms = *this; FusedAttnConfig cfg{}; - // Forward execution: only the forward graph is run, so do not pay for a backward support - // check whose graph this call will never execute. - cfg.check_forward = true; - cfg.check_backward = false; + cfg.is_forward = true; cfg.is_training = params.is_training; cfg.deterministic = false; cfg.cuda_graph = params.cuda_graph; @@ -330,10 +255,6 @@ FusedAttnConfig FusedAttnFwdParams::make_config() const { FusedAttnConfig FusedAttnBwdParams::make_config() const { const FusedAttnBwdParams ¶ms = *this; FusedAttnConfig cfg{}; - // Backward execution: only the backward graph is run. check_forward=false also selects the - // backward key normalization in make_cache_key(). - cfg.check_forward = false; - cfg.check_backward = true; cfg.is_training = true; cfg.deterministic = params.deterministic; 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 1218c7860b..ebc5b3eb07 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -79,18 +79,11 @@ struct FusedAttnConfig { 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. check_forward). Added for convinence purposes and do - // not represent any graph properties. - - // Which directions nvte_get_fused_attn_backend_v2() runs a support check for. The execution - // entry points each run one direction and ask about that one only; a support query leaves the - // defaults and asks about both. Checking a direction that is never executed builds a cuDNN - // graph under a cache key nothing consumes. - // - // check_forward doubles as the direction to build the cuDNN graph for, steering - // make_cache_key() normalization, so it must stay true wherever a forward graph is built. - bool check_forward = true; - bool check_backward = true; + // Filled by derive() or set by caller (i.e. is_forward). Added for convinence purposes and do not + // represent any 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; diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index ebabe79844..b3b9922abf 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -300,17 +300,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi std::to_string(static_cast(qkv_format)) + "."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg.check_forward) { - 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; - } + 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 (cfg.is_training && cfg.check_backward) { + if (cfg.is_training && !cfg.is_forward) { std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); if (!bwd_reason.empty()) { - set_message(message, NVTE_FUSED_ATTN_BWD_REJECT_PREFIX + std::move(bwd_reason)); + set_message(message, std::move(bwd_reason)); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } @@ -327,17 +325,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg.check_forward) { - 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; - } + std::string fwd_reason = is_supported_f16_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (cfg.is_training && cfg.check_backward) { + if (cfg.is_training && !cfg.is_forward) { std::string bwd_reason = is_supported_f16_bwd(cfg, handle); if (!bwd_reason.empty()) { - set_message(message, NVTE_FUSED_ATTN_BWD_REJECT_PREFIX + std::move(bwd_reason)); + set_message(message, std::move(bwd_reason)); 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 22ac9481e2..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 @@ -155,32 +155,23 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // 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 graph_cache::SingleFlight sdpa_f16_fprop_cache_sf; + static std::mutex sdpa_f16_fprop_cache_mutex; - // Get plan from cache if available; otherwise build it exactly once across - // threads (single-flight), so concurrent misses of the same key don't each - // compile and discard an identical graph. + // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - auto &sf = sdpa_f16_fprop_cache_sf; + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; { - std::unique_lock lock(sf.mutex); - // Wait until the graph is cached, or no other thread is building this key. - sf.cv.wait(lock, [&] { - return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; - }); + std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); auto it = cache.find(descriptor); - if (it != cache.end()) { - graph_and_tensors cached_graph = it->second; // copy under the lock - lock.unlock(); - graph_cache_debug::record_cache_lookup("fwd", /*hit=*/true, cfg, descriptor.device_id); - return cached_graph; - } - // Claim the build for this key, so a concurrent miss waits instead of - // compiling an identical graph. No claim means no waiting. - if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); + 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; } - graph_cache_debug::record_cache_lookup("fwd", /*hit=*/false, cfg, descriptor.device_id); - graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -463,11 +454,10 @@ 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); graph_cache_debug::record_build("fwd"); - // Insert our graph. With single-flight we are normally the only builder - // for this key; insert() still tolerates a pre-existing entry and returns - // it. claim_guard releases the build claim and wakes waiters on return. + // 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(sf.mutex); + std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -720,32 +710,23 @@ void fused_attn_arbitrary_seqlen_bwd_impl( using CacheType = std::map; static CacheType sdpa_f16_bprop_cache; - static graph_cache::SingleFlight sdpa_f16_bprop_cache_sf; + static std::mutex sdpa_f16_bprop_cache_mutex; - // Get plan from cache if available; otherwise build it exactly once across - // threads (single-flight), so concurrent misses of the same key don't each - // compile and discard an identical graph. + // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { - auto &sf = sdpa_f16_bprop_cache_sf; + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; { - std::unique_lock lock(sf.mutex); - // Wait until the graph is cached, or no other thread is building this key. - sf.cv.wait(lock, [&] { - return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; - }); + std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); auto it = cache.find(descriptor); - if (it != cache.end()) { - graph_and_tensors cached_graph = it->second; // copy under the lock - lock.unlock(); - graph_cache_debug::record_cache_lookup("bwd", /*hit=*/true, cfg, descriptor.device_id); - return cached_graph; - } - // Claim the build for this key, so a concurrent miss waits instead of - // compiling an identical graph. No claim means no waiting. - if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); + 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; } - graph_cache_debug::record_cache_lookup("bwd", /*hit=*/false, cfg, descriptor.device_id); - graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -999,11 +980,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, offset_s_tuple, dropout_tuple); graph_cache_debug::record_build("bwd"); - // Insert our graph. With single-flight we are normally the only builder - // for this key; insert() still tolerates a pre-existing entry and returns - // it. claim_guard releases the build claim and wakes waiters on return. + // 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(sf.mutex); + std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -1363,11 +1343,8 @@ 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.check_forward = true; + graph_cfg.is_forward = true; graph_cfg.derive(); - // Attribute the graph this builds to the support probe, not to a real - // execution: it may be for a config that never runs. - graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { @@ -1391,11 +1368,8 @@ 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.check_forward = false; + graph_cfg.is_forward = false; graph_cfg.derive(); - // Attribute the graph this builds to the support probe, not to a real - // execution: it may be for a config that never runs. - graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 60699cef0a..9449f74206 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -130,32 +130,23 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de // 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 graph_cache::SingleFlight sdpa_fp8_fprop_cache_sf; + static std::mutex sdpa_fp8_fprop_cache_mutex; - // Get plan from cache if available; otherwise build it exactly once across - // threads (single-flight), so concurrent misses of the same key don't each - // compile and discard an identical graph. + // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - auto& sf = sdpa_fp8_fprop_cache_sf; + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; { - std::unique_lock lock(sf.mutex); - // Wait until the graph is cached, or no other thread is building this key. - sf.cv.wait(lock, [&] { - return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; - }); + std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); auto it = cache.find(descriptor); - if (it != cache.end()) { - graph_and_tensors cached_graph = it->second; // copy under the lock - lock.unlock(); - graph_cache_debug::record_cache_lookup("fwd", /*hit=*/true, cfg, descriptor.device_id); - return cached_graph; - } - // Claim the build for this key, so a concurrent miss waits instead of - // compiling an identical graph. No claim means no waiting. - if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); + 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; } - graph_cache_debug::record_cache_lookup("fwd", /*hit=*/false, cfg, descriptor.device_id); - graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -420,11 +411,10 @@ 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); graph_cache_debug::record_build("fwd"); - // Insert our graph. With single-flight we are normally the only builder - // for this key; insert() still tolerates a pre-existing entry and returns - // it. claim_guard releases the build claim and wakes waiters on return. + // 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(sf.mutex); + std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -636,32 +626,23 @@ void fused_attn_fp8_bwd_impl( using CacheType = std::map; static CacheType sdpa_fp8_bprop_cache; - static graph_cache::SingleFlight sdpa_fp8_bprop_cache_sf; + static std::mutex sdpa_fp8_bprop_cache_mutex; - // Get plan from cache if available; otherwise build it exactly once across - // threads (single-flight), so concurrent misses of the same key don't each - // compile and discard an identical graph. + // Get plan from cache if cache is available, otherwise create one auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { - auto& sf = sdpa_fp8_bprop_cache_sf; + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; { - std::unique_lock lock(sf.mutex); - // Wait until the graph is cached, or no other thread is building this key. - sf.cv.wait(lock, [&] { - return cache.count(descriptor) != 0 || sf.in_progress.count(descriptor) == 0; - }); + std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); auto it = cache.find(descriptor); - if (it != cache.end()) { - graph_and_tensors cached_graph = it->second; // copy under the lock - lock.unlock(); - graph_cache_debug::record_cache_lookup("bwd", /*hit=*/true, cfg, descriptor.device_id); - return cached_graph; - } - // Claim the build for this key, so a concurrent miss waits instead of - // compiling an identical graph. No claim means no waiting. - if (graph_cache::single_flight_enabled()) sf.in_progress.insert(descriptor); + 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; } - graph_cache_debug::record_cache_lookup("bwd", /*hit=*/false, cfg, descriptor.device_id); - graph_cache::ClaimGuard claim_guard{sf, descriptor}; // otherwise, build the op_graph and the plan. Then update cache auto mha_graph = std::make_shared(); @@ -1056,11 +1037,10 @@ 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); graph_cache_debug::record_build("bwd"); - // Insert our graph. With single-flight we are normally the only builder - // for this key; insert() still tolerates a pre-existing entry and returns - // it. claim_guard releases the build claim and wakes waiters on return. + // 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(sf.mutex); + std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); auto inserted = cache.insert({descriptor, return_tuple}); return inserted.first->second; } @@ -1405,11 +1385,8 @@ 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.check_forward = true; + graph_cfg.is_forward = true; graph_cfg.derive(); - // Attribute the graph this builds to the support probe, not to a real - // execution: it may be for a config that never runs. - graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { @@ -1434,11 +1411,8 @@ 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.check_forward = false; + graph_cfg.is_forward = false; graph_cfg.derive(); - // Attribute the graph this builds to the support probe, not to a real - // execution: it may be for a config that never runs. - graph_cache_debug::ScopedProbe probe; size_t workspace_size = 0; try { diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 3961c0871c..033edf2ead 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -7,42 +7,22 @@ // ============================================================================ // Fused-attention graph cache diagnostics. // -// Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG. Two verbosity levels: -// =1 : low volume. Cache event counters, per-build BUILD lines, and the -// end-of-run SUMMARY (aggregate + per-thread) and stage timings. This is -// enough to diagnose redundant rebuilds and profile build cost. -// =2 : high volume (trace). Additionally emits a per-lookup HIT/MISS line with -// the full shorthand config and a per-execution EXEC line. Use only when -// you need to see *which* shapes are hitting/missing -- these fire on -// every cache lookup and execution, so at suite scale they add I/O and -// serialize threads on the stderr lock (perturbing the build timings). -// -// An optional ":" suffix picks which processes emit, defaulting to rank 0 -// so that output does not scale with the world size: "1:all" for every rank, -// "2:0,3" for a specific set. See `rank_selected` for when overriding pays off. +// 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 #include #include -#include -#include -#include -#include -#include - -#include -#include #include "config_and_params.h" @@ -50,129 +30,23 @@ namespace transformer_engine { namespace fused_attn { namespace graph_cache_debug { -// Rank of this process as reported by the launcher, or -1 when there is no -// launcher (a single-process run). First variable that is set wins. -inline int launcher_rank() { - static const int rank = []() -> int { - for (const char *var : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { - const char *v = std::getenv(var); - if (v != nullptr && v[0] != '\0') return std::atoi(v); - } - return -1; - }(); - return rank; -} - -// Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=default, -// 2=trace). Single read at startup, cached. Negligible overhead when unset. -inline int debug_level() { - static const int lvl = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; - const int v = std::atoi(e); // stops at the optional ":" suffix - return v > 0 ? v : 1; // any non-empty, non-"0" value enables at least level 1 - }(); - return lvl; -} - -// Whether this process emits diagnostics. Every rank writes to the same stderr, -// so emitting from all of them multiplies the volume by the world size -- and -// under data/tensor parallelism the ranks are running identical shapes, so the -// copies say the same thing. Hence rank 0 only by default. -// -// Context parallelism is the case worth overriding for: the ranks run different -// subsets of the per-step regimes (under p2p, rank 0 never sees the lower-triangle -// config that the last rank does), so their build counts genuinely differ. -// Select with the ":" suffix, e.g. "1:all" or "2:0,3". -inline bool rank_selected() { - static const bool selected = [] { - const int rank = launcher_rank(); - if (rank < 0) return true; // sole process, nothing to filter +// 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"); - const char *sep = (e != nullptr) ? std::strchr(e, ':') : nullptr; - if (sep == nullptr) return rank == 0; - const std::string list(sep + 1); - if (list == "all") return true; - for (size_t pos = 0; pos <= list.size();) { - const size_t comma = list.find(',', pos); - const std::string tok = - list.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); - if (!tok.empty() && std::atoi(tok.c_str()) == rank) return true; - if (comma == std::string::npos) break; - pos = comma + 1; - } - return false; - }(); - return selected; -} - -// Diagnostics are on at level >= 1, and only for the selected ranks. Unselected -// ranks skip the counters too, so they pay nothing beyond this check. -inline bool enabled() { return debug_level() >= 1 && rank_selected(); } - -// Per-lookup / per-exec trace lines are gated behind level >= 2. -inline bool trace_enabled() { return debug_level() >= 2; } - -// Identifies the emitting process. Distributed PyTorch runs one process per rank -// and they all share this stderr, so without this every line would be ambiguous -// (thread ids restart at 0 in each process). Rank comes from the launcher, if any. -inline const std::string &process_tag() { - static const std::string *tag = [] { - auto *s = new std::string("pid=" + std::to_string(static_cast(::getpid()))); - if (launcher_rank() >= 0) *s += " rank=" + std::to_string(launcher_rank()); - return s; + return e != nullptr && e[0] != '\0' && e[0] != '0'; }(); - return *tag; + return on; } -// More readable, shorter thread IDs (0, 1, 2, ...). These are assignment order, -// not identity: tid=0 is whichever thread touched this cache first. The one-shot -// THREAD line below maps them to OS thread ids for correlating with nsys/gdb. +// 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; } -// OS-level thread id, as reported by nsys/gdb/`top -H`. Via syscall rather than -// gettid() so this does not require glibc >= 2.30. -inline int64_t os_thread_id() { return static_cast(::syscall(SYS_gettid)); } - -// True while this thread is inside a backend-support probe (`is_supported_*`), -// which builds a graph speculatively just to answer "is this config supported?". -// Such a build may never be executed -- notably the context-parallel per-step -// probe, which checks regimes this rank never runs -- so it is counted apart -// from builds triggered by a real execution. -inline bool &tl_in_probe() { - static thread_local bool v = false; - return v; -} - -// Marks the calling scope as a support probe. Saves/restores rather than -// clearing, so it stays correct if probes ever nest. -struct ScopedProbe { - bool prev; - ScopedProbe() : prev(tl_in_probe()) { tl_in_probe() = true; } - ~ScopedProbe() { tl_in_probe() = prev; } -}; - -inline const char *src_tag() { return tl_in_probe() ? "probe" : "exec"; } - -// Milliseconds since the first diagnostic event. Used to correlate build -// start/end across threads and detect whether same-shape builds on different -// devices overlap in wall-clock time. -inline double now_ms() { - static const auto t0 = std::chrono::steady_clock::now(); - return std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); -} - -// Per-thread build-start timestamp, set at the miss that triggers a build and -// read when the build completes, to report each build's duration. -inline double &tl_build_start_ms() { - static thread_local double v = 0.0; - return v; -} - // Registered at first use. On process exit, prints overall event counters and // graph build timings. inline void register_summary_once(); @@ -189,9 +63,6 @@ inline void register_summary_once(); struct EventCounters { std::atomic built{0}; - // Subset of `built` that a support probe triggered. A large probe share means - // time is going into graphs that may never run. - std::atomic built_probe{0}; std::atomic exec{0}; std::atomic hit{0}; std::atomic miss{0}; @@ -203,110 +74,27 @@ inline EventCounters &counters(bool is_fwd) { return is_fwd ? fwd : bwd; } -// Per-thread counters, so the summary can break down build/exec/hit/miss by -// thread. In the single-process context-parallel case each device is driven by -// its own thread, so this reveals which thread built/executed what. -struct ThreadCounters { - unsigned tid = 0; - EventCounters fwd; - EventCounters bwd; -}; - -// The registry and its mutex are heap-allocated and deliberately never freed. -// Function-local static destructors and atexit handlers run as a single sequence, -// in reverse order of construction/registration. This registry is built lazily, so -// it can be constructed *after* the summary handler is registered -- in which case -// it would be destroyed *before* that handler runs, leaving the handler to lock a -// destroyed mutex and walk a destroyed vector. Leaking removes the ordering -// question rather than reasoning about it, and the cost is bounded: one mutex and -// one vector for the process, reclaimed by the OS at exit anyway. -inline std::mutex &thread_registry_mutex() { - static std::mutex *m = new std::mutex(); - return *m; -} -inline std::vector &thread_registry() { - static std::vector *v = new std::vector(); - return *v; -} - -// This thread's counter block, leaked for a related but distinct reason: a worker -// thread can exit long before the process does, while the registry keeps a pointer -// to its block for the end-of-run summary. Tying the block's lifetime to the -// thread would leave that pointer dangling. One small struct per thread. -inline ThreadCounters &thread_counters() { - static thread_local ThreadCounters *tc = [] { - auto *p = new ThreadCounters(); - p->tid = thread_seq_id(); - { - std::lock_guard lock(thread_registry_mutex()); - thread_registry().push_back(p); - } - // One line per thread, mapping the short id to something nsys/gdb can match. - std::fprintf(stderr, "[FUSED-ATTN-CACHE] %s | THREAD | tid=%-3u os_tid=%" PRId64 "\n", - process_tag().c_str(), p->tid, os_thread_id()); - std::fflush(stderr); - return p; - }(); - return *tc; -} - -inline EventCounters &thread_counters(bool is_fwd) { - ThreadCounters &tc = thread_counters(); - return is_fwd ? tc.fwd : tc.bwd; -} - -// Format one counter block (aggregate or a single thread's) as one line. -// `tid_field` is the whole thread column, e.g. "tid=3"; the aggregate row passes -// "tid=all" so that it cannot be misread as thread 0's row. -inline std::string format_counter_line(const char *event, const char *tid_field, - const EventCounters &f, const EventCounters &b, - const char *extra) { - char buf[640]; - std::snprintf(buf, sizeof(buf), - "[FUSED-ATTN-CACHE] %s | %-11s | %-7s | fwd miss=%4" PRIu64 ", hit=%4" PRIu64 - ", built=%4" PRIu64 " (for probe %4" PRIu64 "), exec=%4" PRIu64 " | bwd miss=%4" PRIu64 - ", hit=%4" PRIu64 ", built=%4" PRIu64 " (for probe %4" PRIu64 "), exec=%4" PRIu64 "%s\n", - process_tag().c_str(), event, tid_field, f.miss.load(std::memory_order_relaxed), - f.hit.load(std::memory_order_relaxed), f.built.load(std::memory_order_relaxed), - f.built_probe.load(std::memory_order_relaxed), - f.exec.load(std::memory_order_relaxed), b.miss.load(std::memory_order_relaxed), - b.hit.load(std::memory_order_relaxed), b.built.load(std::memory_order_relaxed), - b.built_probe.load(std::memory_order_relaxed), - b.exec.load(std::memory_order_relaxed), extra); - return std::string(buf); -} - -inline void print_counter_block(const char *event, const char *tid_field, const EventCounters &f, - const EventCounters &b, const char *extra = "") { - const std::string line = format_counter_line(event, tid_field, f, b, extra); - std::fputs(line.c_str(), stderr); +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=%" 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); } -inline void print_counters(const char *event, const char *extra = "") { - char tid_field[16]; - std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); - print_counter_block(event, tid_field, counters(/*is_fwd=*/true), counters(/*is_fwd=*/false), - extra); -} - 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); - thread_counters(is_fwd).built.fetch_add(1, std::memory_order_relaxed); - if (tl_in_probe()) { - counters(is_fwd).built_probe.fetch_add(1, std::memory_order_relaxed); - thread_counters(is_fwd).built_probe.fetch_add(1, std::memory_order_relaxed); - } - // Report build completion time and this build's wall-clock duration so we can - // tell whether same-shape builds on different devices overlap. - const double t_end = now_ms(); - char extra[80]; - std::snprintf(extra, sizeof(extra), " | src=%-5s t=%.1f dur=%.1f ms", src_tag(), t_end, - t_end - tl_build_start_ms()); - print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD", extra); + print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD"); } inline void record_exec(const char *pass) { @@ -314,36 +102,17 @@ inline void record_exec(const char *pass) { register_summary_once(); const bool is_fwd = std::strcmp(pass, "fwd") == 0; counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); - thread_counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); - // The per-exec line fires on every execution; keep it out of the level-1 path. - if (!trace_enabled()) return; - char extra[32]; - std::snprintf(extra, sizeof(extra), " | t=%.1f", now_ms()); - print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC", extra); + print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC"); } -// `device_key` is the cache-scope discriminator from make_cache_key(), not a device ordinal: -// it is the packed (SM arch, SM count) when devices share plans, else the device id. -inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c, - int device_key) { +inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { if (!enabled()) return; register_summary_once(); - const bool is_fwd = std::strcmp(pass, "fwd") == 0; - EventCounters &pc = counters(is_fwd); + EventCounters &pc = counters(std::strcmp(pass, "fwd") == 0); (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); - EventCounters &tpc = thread_counters(is_fwd); - (hit ? tpc.hit : tpc.miss).fetch_add(1, std::memory_order_relaxed); - const double t = now_ms(); - // A miss triggers a build right after this call; stamp the build start so the - // subsequent BUILD line can report duration. Do this even at level 1. - if (!hit) tl_build_start_ms() = t; - // The per-lookup config dump is the highest-volume line (one per cache probe); - // keep it out of the level-1 path and off the stderr lock unless tracing. - if (!trace_enabled()) return; std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %s | %-3s %-4s | tid=%u devkey=%d t=%.1f src=%-5s | train=%d det=%d cg=%d " - "maxlogit=%d fwd=%d " + "[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 @@ -352,10 +121,9 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi " 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", - process_tag().c_str(), pass, hit ? "HIT" : "MISS", thread_seq_id(), device_key, t, src_tag(), - static_cast(c.is_training), + 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.check_forward), + 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), @@ -377,6 +145,7 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi 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); } // ============================================================================ @@ -392,14 +161,11 @@ struct StageTiming { std::atomic calls{0}; std::atomic time_ns{0}; }; -// Bucketed by pass and by whether a support probe drove the build, so the -// summary can say how much of (notably) `build_plans` was speculative. -constexpr size_t kStageBuckets = 4 * static_cast(BuildStage::kCount); -inline StageTiming &stage_timing(bool is_fwd, bool is_probe, BuildStage s) { +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) * 2u + (is_probe ? 1u : 0u)) * - static_cast(BuildStage::kCount) + - static_cast(s); + const size_t idx = + (is_fwd ? 0u : 1u) * static_cast(BuildStage::kCount) + static_cast(s); return table[idx]; } @@ -407,10 +173,8 @@ struct ScopedBuildTimer { BuildStage stage; bool on; bool is_fwd; - bool is_probe; std::chrono::steady_clock::time_point start; - ScopedBuildTimer(bool is_fwd_, BuildStage s) - : stage(s), on(enabled()), is_fwd(is_fwd_), is_probe(tl_in_probe()) { + 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(); @@ -421,7 +185,7 @@ struct ScopedBuildTimer { static_cast(std::chrono::duration_cast( std::chrono::steady_clock::now() - start) .count()); - StageTiming &t = stage_timing(is_fwd, is_probe, stage); + 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); } @@ -440,51 +204,23 @@ inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; - // Build the whole summary in memory and emit it with a single write, so - // that the blocks of concurrently-exiting processes (one per rank under - // torchrun) stay grouped instead of interleaving line by line. - std::string block; - block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary begin =====\n"; - // Per-thread breakdown (sorted by tid). Useful in the single-process - // context-parallel case where each device runs on its own thread. - { - std::lock_guard lock(thread_registry_mutex()); - std::vector blocks = thread_registry(); - std::sort(blocks.begin(), blocks.end(), - [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); - for (const ThreadCounters *tc : blocks) { - char tid_field[16]; - std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); - block += format_counter_line("SUMMARY-TID", tid_field, tc->fwd, tc->bwd, ""); - } - } - // Totals last, so they read as the sum of the per-thread lines above. - block += format_counter_line("SUMMARY", "tid=all", counters(/*is_fwd=*/true), - counters(/*is_fwd=*/false), ""); + 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 q = 0; q < 2; ++q) { - const bool is_probe = (q == 1); - for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { - const BuildStage s = static_cast(i); - const StageTiming &t = stage_timing(is_fwd, is_probe, 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; - char line[288]; - std::snprintf(line, sizeof(line), - "[FUSED-ATTN-CACHE] %s | %-3s src=%-5s %-22s | calls=%" PRIu64 - " | time=%9.3f ms/call\n", - process_tag().c_str(), pass, is_probe ? "probe" : "exec", kStageNames[i], - n, total_ms / n); - block += line; - } + 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=%" PRIu64 + " | time=%9.1f ms | avg=%9.3f ms/call\n", + pass, kStageNames[i], n, total_ms, total_ms / n); } } - block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary end =====\n"; - std::fwrite(block.data(), 1, block.size(), stderr); std::fflush(stderr); }); return true; @@ -493,74 +229,6 @@ inline void register_summary_once() { } } // namespace graph_cache_debug - -// ============================================================================ -// Single-flight graph cache coordination. -// -// The fused-attention graph caches are process-wide and shared across threads. -// The lock is intentionally released while a graph is compiled so that -// *different* graphs can build in parallel. The downside is that when several -// threads miss the *same* key at the same instant (e.g. the device-worker -// threads of a single-process context-parallel run stepping in lockstep), they -// all compile an identical graph and all but one discard the result at insert -// time -- wasted host-side `build_plans` work. -// -// A single-flight (a.k.a. "thundering herd") guard closes that gap: at most one -// thread compiles a given key while the others wait for it. Distinct keys still -// build concurrently, so the parallel-build win is kept. -// -// Usage in a get_graph path: -// static SingleFlight sf; -// { -// std::unique_lock lock(sf.mutex); -// sf.cv.wait(lock, [&] { -// return cache.count(key) != 0 || sf.in_progress.count(key) == 0; -// }); -// if (auto it = cache.find(key); it != cache.end()) { ...cache hit... } -// sf.in_progress.insert(key); // claim the build -// } -// ClaimGuard guard{sf, key}; // auto-release + notify -// ...build... -// { std::lock_guard lock(sf.mutex); cache.insert({key, graph}); } -// ============================================================================ -namespace graph_cache { - -// Opt-in with NVTE_FUSED_ATTN_CACHE_SINGLE_FLIGHT=1. When off, no thread claims a -// key, so the wait below falls through immediately and concurrent misses of one -// key each compile, with all but one result discarded at insert: wasted host work, -// but no thread ever blocks on another thread's compile. -inline bool single_flight_enabled() { - static const bool on = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_SINGLE_FLIGHT"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; - }(); - return on; -} - -template -struct SingleFlight { - std::mutex mutex; - std::condition_variable cv; - std::set in_progress; -}; - -// RAII: on scope exit, drop this thread's build claim on `key` and wake any -// threads waiting on the same key. Clears the claim even if the build throws, -// so waiters never deadlock (they simply re-elect a builder). -template -struct ClaimGuard { - SingleFlight &sf; - const KeyT &key; - ~ClaimGuard() { - { - std::lock_guard lock(sf.mutex); - sf.in_progress.erase(key); - } - sf.cv.notify_all(); - } -}; - -} // namespace graph_cache } // namespace fused_attn } // 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 0369ed62bf..938fa1747e 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -413,17 +413,6 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, const void *buf, size_t size_in_bytes); -/*! \brief Prefix on the ``nvte_get_fused_attn_backend_v2`` diagnostic when a configuration is - * rejected solely by the backward-pass support check. - * - * Its presence tells the caller that the forward pass is supported and only the backward pass - * is not, so re-querying the same configuration with ``is_training = false`` may succeed. Its - * absence means the forward pass itself was rejected, for which dropping to inference cannot - * help. Mirrored on the Python side as - * ``transformer_engine.pytorch.attention.dot_product_attention.utils.FUSED_ATTN_BWD_REJECT_PREFIX``. - */ -#define NVTE_FUSED_ATTN_BWD_REJECT_PREFIX "[backward] " - /*! \brief Get fused-attention backend based on user configuration. * * This function passes the user configuration to cuDNN frontend, runs its support checks, @@ -433,9 +422,7 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, * \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 explaining why there is no support, - * prefixed with ``NVTE_FUSED_ATTN_BWD_REJECT_PREFIX`` when only the - * backward pass is unsupported. + * 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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/__init__.py b/transformer_engine/pytorch/attention/dot_product_attention/__init__.py index bd54ca3134..941f94f105 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/__init__.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/__init__.py @@ -4,18 +4,6 @@ """Python interface for dot product attention""" -from .dot_product_attention import ( - DotProductAttention, - BackendSelectionProbe, - DryRunResult, - dry_run_backend_selection, - _attention_backends, -) +from .dot_product_attention import DotProductAttention, _attention_backends -__all__ = [ - "DotProductAttention", - "BackendSelectionProbe", - "DryRunResult", - "dry_run_backend_selection", - "_attention_backends", -] +__all__ = ["DotProductAttention", "_attention_backends"] 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 cd9555da52..93296e62dd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4946,8 +4946,7 @@ def cp_per_step_configs( 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, t_q, t_kv, window=None): - w_left, w_right = window if window is not None else (window_left, window_right) + def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): return { "attn_mask_type": mask, "max_seqlen_q": s_q, @@ -4956,8 +4955,8 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv, window=None): "num_tokens_kv": t_kv, "num_attn_heads": heads, "num_gqa_groups": gqa, - "window_size_left": w_left, - "window_size_right": w_right, + "window_size_left": window_left, + "window_size_right": window_right, "bottom_right_diagonal": bottom_right, } @@ -4983,16 +4982,20 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv, window=None): 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 - # Each step narrows max_seqlen_*, but the token counts it dispatches with are the - # rank's full Q tokens and the all-gathered KV tokens, unchanged across steps. - # Scaling them per step would key the probe's graph differently from the one the - # step looks up, and rebuild every graph this probes at execution time. - t_q = num_tokens_q - t_kv = num_tokens_kv * cp_size + t_q = num_tokens_q // 2 # 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, t_q, t_kv) + config( + mask, + s_q, + s_kv, + num_heads, + num_gqa_groups, + br, + t_q, + num_tokens_kv * cp_size * s_kv // max_seqlen_kv if max_seqlen_kv else 0, + ) for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) ] 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 2117c528c7..3c079cc5d5 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 @@ -3,11 +3,9 @@ # See LICENSE for license information. """Attention.""" -from contextlib import contextmanager, nullcontext -import dataclasses +from contextlib import nullcontext import math import os -import threading from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings import logging @@ -74,122 +72,9 @@ "use_fused_attention": None, "fused_attention_backend": None, "use_unfused_attention": None, - "available_backends": None, - "fused_attn_reject_reason": None, "backend_selection_requires_update": False, } -# Dry-run backend selection: see dry_run_backend_selection(). -_dpa_dry_run = threading.local() - - -class _DryRunComplete(Exception): - """Unwinds the forward pass once backend selection is known.""" - - -@dataclasses.dataclass -class BackendSelectionProbe: - """Backend selection at a single `DotProductAttention` site. - - Attributes - ---------- - attention_params : AttentionParams - The parameters the module actually resolved, i.e. the exact cache key the real run - would use. - available_backends : List[bool] - [flash, fused, unfused] support for this configuration. - fused_attention_backend : Optional[FusedAttnBackend] - The selected `FusedAttention` sub-backend, or `None`. - fused_attn_reject_reason : Optional[str] - Why `FusedAttention` was ruled out, or `None`. A reason starting with - `FUSED_ATTN_BWD_REJECT_PREFIX` means only the backward pass is unsupported. - """ - - attention_params: Optional["dpa_utils.AttentionParams"] = None - available_backends: Optional[List[bool]] = None - fused_attention_backend: Optional[Any] = None - fused_attn_reject_reason: Optional[str] = None - - -@dataclasses.dataclass -class DryRunResult: - """Backend selection across every `DotProductAttention` site a dry run reached. - - A module runs only if all of its attention sites are supported, so the `*_supported` - properties require every recorded site to support the backend. `probes` holds the - per-site detail, in the order the sites were reached. - """ - - stop_after: int = 1 - probes: List[BackendSelectionProbe] = dataclasses.field(default_factory=list) - - def _all_support(self, index: int) -> bool: - return bool(self.probes) and all(p.available_backends[index] for p in self.probes) - - @property - def flash_supported(self) -> bool: - """Whether `FlashAttention` supports every attention site reached.""" - return self._all_support(0) - - @property - def fused_supported(self) -> bool: - """Whether `FusedAttention` supports every attention site reached.""" - return self._all_support(1) - - @property - def unfused_supported(self) -> bool: - """Whether `UnfusedDotProductAttention` supports every attention site reached.""" - return self._all_support(2) - - @property - def fused_attn_reject_reason(self) -> Optional[str]: - """Why `FusedAttention` was ruled out, from the first site that ruled it out.""" - return next( - (p.fused_attn_reject_reason for p in self.probes if p.fused_attn_reject_reason), - None, - ) - - -@contextmanager -def dry_run_backend_selection(stop_after: int = 1): - """Resolve attention backends for a module without running attention. - - Call a module as usual inside this context. `DotProductAttention.forward` resolves its - configuration exactly as a real run would, records which backends support it, and then - unwinds the forward pass before any attention is executed. This avoids having to - predict the configuration a module will produce, e.g. the `qkv_layout` that - `MultiheadAttention` derives from its packed projection output. - - Because the resolved configuration is identical to the real run's, the backend query - populates the same cache entries the real run will hit. - - Parameters - ---------- - stop_after : int, default = 1 - Unwind after this many attention sites have been recorded. Modules with several - attention sites, e.g. a `TransformerLayer` with `layer_type="decoder"`, need a - higher value to reach the later ones. Note that the sites before the last one - execute for real, since that is the only way to arrive at what follows them. - - .. code-block:: python - - with dry_run_backend_selection() as dry_run: - model(hidden_states, attn_mask_type="causal") - if not dry_run.fused_supported: - print(dry_run.fused_attn_reject_reason) - """ - assert stop_after >= 1, "stop_after must be at least 1" - result = DryRunResult(stop_after=stop_after) - previous = getattr(_dpa_dry_run, "result", None) - _dpa_dry_run.result = result - try: - yield result - except _DryRunComplete: - pass - finally: - _dpa_dry_run.result = previous - - _alibi_cache = { "_num_heads": None, "_alibi_slopes": None, @@ -2156,9 +2041,6 @@ def forward( use_flash_attention = False use_fused_attention = False use_unfused_attention = True - fused_attention_backend = None - available_backends = [False, False, True] - fused_attn_reject_reason = None else: if ( _attention_backends["attention_params"] is None @@ -2173,8 +2055,7 @@ def forward( use_fused_attention, fused_attention_backend, use_unfused_attention, - available_backends, - fused_attn_reject_reason, + _, ) = dpa_utils.get_attention_backend(attention_params) # Set global _attention_backends var using return value # from get_attention_backend() @@ -2183,8 +2064,6 @@ def forward( _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["available_backends"] = available_backends - _attention_backends["fused_attn_reject_reason"] = fused_attn_reject_reason _attention_backends["backend_selection_requires_update"] = False if use_flash_attention: self.logger.info( @@ -2204,21 +2083,6 @@ def forward( use_fused_attention = _attention_backends["use_fused_attention"] fused_attention_backend = _attention_backends["fused_attention_backend"] use_unfused_attention = _attention_backends["use_unfused_attention"] - available_backends = _attention_backends["available_backends"] - fused_attn_reject_reason = _attention_backends["fused_attn_reject_reason"] - - dry_run = getattr(_dpa_dry_run, "result", None) - if dry_run is not None: - dry_run.probes.append( - BackendSelectionProbe( - attention_params=attention_params, - available_backends=available_backends, - fused_attention_backend=fused_attention_backend, - fused_attn_reject_reason=fused_attn_reject_reason, - ) - ) - if len(dry_run.probes) >= dry_run.stop_after: - raise _DryRunComplete # 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 f37bdacefc..eb10f0c03d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -73,11 +73,6 @@ _cu_seqlens_cache = {} -# Mirrors NVTE_FUSED_ATTN_BWD_REJECT_PREFIX in common/include/transformer_engine/fused_attn.h. -# A rejection reason carrying this prefix means only the backward pass is unsupported, so the -# same config may be supported with is_training=False. Keep the two definitions in sync. -FUSED_ATTN_BWD_REJECT_PREFIX = "[backward] " - class AttentionLogging: """ @@ -471,11 +466,6 @@ def get_attention_backend( available_backends : List[bool] 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]. - fused_attention_reject_reason : Optional[str] - Why `FusedAttention` was ruled out, or `None` if it was selected or was never queried - (e.g. disabled via `NVTE_FUSED_ATTN=0`). A reason starting with - `FUSED_ATTN_BWD_REJECT_PREFIX` means only the backward pass is unsupported, so the same - config may be supported with `is_training=False`. """ # 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 @@ -1541,7 +1531,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # Filter: cuDNN support fused_attention_backend = None - fused_attention_reject_reason = None if use_fused_attention: 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" @@ -1653,11 +1642,10 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt window_size_right=step_config["window_size_right"], bottom_right_diagonal=step_config["bottom_right_diagonal"], ) - if fu_core_attention_bias_shape is not None: - 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 + 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 # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. @@ -1674,7 +1662,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_fused_attention = False fused_attention_backend = None - fused_attention_reject_reason = reject_message break if ( @@ -1735,11 +1722,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): - reason = "determinism with FP8 is not supported on arch < sm90 or cuDNN < 9.19.0" - logger.debug("Disabling FusedAttention for %s", reason) + logger.debug( + "Disabling FusedAttention for determinism reasons with FP8 on arch < sm90 or cuDNN" + " < 9.19.0" + ) use_fused_attention = False fused_attention_backend = None - fused_attention_reject_reason = FUSED_ATTN_BWD_REJECT_PREFIX + reason if ( fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] and is_training @@ -1749,11 +1737,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt or cudnn_version < (8, 9, 5) ) ): - reason = "determinism with post_scale_bias is not supported" - logger.debug("Disabling FusedAttention for %s", reason) + logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False fused_attention_backend = None - fused_attention_reject_reason = FUSED_ATTN_BWD_REJECT_PREFIX + reason # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 @@ -1863,7 +1849,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fused_attention_backend, use_unfused_attention, available_backends, - fused_attention_reject_reason, ) From 734aca7a912426645ad23153d8892df27536fc46 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:24:33 -0700 Subject: [PATCH 72/88] fix probe/exec bias drift in keys Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../pytorch/attention/dot_product_attention/utils.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index eb10f0c03d..01eb4a960c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1642,10 +1642,11 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt 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 + if fu_core_attention_bias_shape is not None: + 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 # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. From d2f7774d2316b6d6faf3dc0bbd556e4b4a8a74d1 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:25:37 -0700 Subject: [PATCH 73/88] fix probe/exec cp drift in keys Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../dot_product_attention/context_parallel.py | 25 ++++++++----------- 1 file changed, 11 insertions(+), 14 deletions(-) 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 93296e62dd..cd9555da52 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4946,7 +4946,8 @@ def cp_per_step_configs( 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, t_q, t_kv): + def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv, window=None): + w_left, w_right = window if window is not None else (window_left, window_right) return { "attn_mask_type": mask, "max_seqlen_q": s_q, @@ -4955,8 +4956,8 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): "num_tokens_kv": t_kv, "num_attn_heads": heads, "num_gqa_groups": gqa, - "window_size_left": window_left, - "window_size_right": window_right, + "window_size_left": w_left, + "window_size_right": w_right, "bottom_right_diagonal": bottom_right, } @@ -4982,20 +4983,16 @@ def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv): 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 - t_q = num_tokens_q // 2 + # Each step narrows max_seqlen_*, but the token counts it dispatches with are the + # rank's full Q tokens and the all-gathered KV tokens, unchanged across steps. + # Scaling them per step would key the probe's graph differently from the one the + # step looks up, and rebuild every graph this probes at execution time. + t_q = num_tokens_q + t_kv = num_tokens_kv * cp_size # 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, - t_q, - num_tokens_kv * cp_size * s_kv // max_seqlen_kv if max_seqlen_kv else 0, - ) + config(mask, s_q, s_kv, num_heads, num_gqa_groups, br, t_q, t_kv) for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) ] From 6201d048289615a336083b0253bebe059b280b2d Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:26:01 -0700 Subject: [PATCH 74/88] fix MHA init to avoid probe/exec drift Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index ca93fecc9e..a800df578c 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2300,6 +2300,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attention_type="self", qkv_weight_interleaved=True, qkv_format=qkv_format, + window_size=config.window_size, + softmax_type=config.softmax_type, ).to(dtype=dtype, device="cuda") if not is_training: mha = mha.eval() From 157a975fd8769abb0518cc478149b3d0166f6ca7 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:56:25 -0700 Subject: [PATCH 75/88] WIP: graph cache and restructuring of impl 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 | 33 +- .../common/fused_attn/fused_attn.cpp | 35 +- .../fused_attn_f16_arbitrary_seqlen.cu | 1618 +++++++------- .../common/fused_attn/fused_attn_fp8.cu | 1853 +++++++++-------- .../common/fused_attn/graph_cache.h | 242 +++ .../common/fused_attn/graph_cache_debug.h | 422 +++- transformer_engine/common/fused_attn/utils.h | 1 - 8 files changed, 2437 insertions(+), 1787 deletions(-) create mode 100644 transformer_engine/common/fused_attn/graph_cache.h diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index ca4214dac3..91b4445617 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -91,6 +91,8 @@ void FusedAttnConfig::derive() { max_pages_per_seq_v = 1; } } + + is_derived = true; } FusedAttnConfig FusedAttnConfig::make_cache_key() const { @@ -121,7 +123,12 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { } cache_cfg.num_tokens_q = 0; cache_cfg.num_tokens_kv = 0; - const bool bucket_batch = !is_forward || !cache_cfg.uses_cu_seqlens_directly; + // The forward graph keeps the true batch size when it takes the user's cu_seqlens + // directly, since cuDNN reads those [actual_b+1] buffers itself; the backward graph + // converts them and so always buckets. The key has to follow whichever the graph does, + // or it would name a batch size the graph was not built with. See + // derive_f16_bwd_graph_inputs. + const bool bucket_batch = !check_for_forward_support || !cache_cfg.uses_cu_seqlens_directly; if (bucket_batch) { cache_cfg.batch_size = cache_cfg.bucketed_batch_size; } @@ -133,7 +140,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { // 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 (check_for_forward_support) { cache_cfg.do_dtype = kNVTEBFloat16; cache_cfg.dqkv_dtype = kNVTEBFloat16; cache_cfg.do_format = NVTE_QKV_Format_NOT_SET; @@ -150,7 +157,10 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { FusedAttnConfig FusedAttnFwdParams::make_config() const { const FusedAttnFwdParams ¶ms = *this; FusedAttnConfig cfg{}; - cfg.is_forward = true; + // Forward execution: only the forward graph is run, so do not pay for a backward support + // check whose graph this call will never execute. + cfg.check_for_forward_support = true; + cfg.check_for_backward_support = false; cfg.is_training = params.is_training; cfg.deterministic = false; cfg.cuda_graph = params.cuda_graph; @@ -255,6 +265,10 @@ FusedAttnConfig FusedAttnFwdParams::make_config() const { FusedAttnConfig FusedAttnBwdParams::make_config() const { const FusedAttnBwdParams ¶ms = *this; FusedAttnConfig cfg{}; + // Backward execution: only the backward graph is run. check_for_forward_support=false also + // selects the backward key normalization in make_cache_key(). + cfg.check_for_forward_support = false; + cfg.check_for_backward_support = true; cfg.is_training = true; cfg.deterministic = params.deterministic; 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 ebc5b3eb07..5d04499409 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -79,11 +79,16 @@ struct FusedAttnConfig { 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. - - // Direction to build the cuDNN graph for; steers make_cache_key() normalization. - bool is_forward = false; + // Filled by derive() or set by caller (i.e. check_for_forward_support). Added for convinence + // purposes and do not represent any graph properties. + bool check_for_forward_support = true; + bool check_for_backward_support = true; + // Whether derive() has run, i.e. whether the fields below hold anything. Every consumer of a + // derived field needs them filled -- an unfilled config yields a graph with the wrong shapes + // and a cache key that collides with unrelated configs, neither of which announces itself -- + // so this exists to let those consumers assert rather than trust. Not a cached-result marker: + // derive() recomputes unconditionally, so a config whose inputs change can simply be re-derived. + bool is_derived = 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; @@ -175,7 +180,10 @@ struct FusedAttnConfig { } // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields - // that have been set by the caller. + // that have been set by the caller. Call once, after the last input field is set and before + // the config reaches a graph build, a cache lookup, or a support query -- all of which read + // derived fields. nvte_get_fused_attn_backend_v2() is where that happens for every config + // that enters this library, so nothing downstream of it needs to derive again. void derive(); // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. @@ -185,6 +193,19 @@ struct FusedAttnConfig { FusedAttnConfig make_cache_key() const; }; +// Assert that `cfg` has been through derive(), for code about to read a derived field. Worth +// asserting rather than assuming because the failure is silent: an unset bucketed_batch_size or +// q_format reads as zero, which is a legal value that yields a graph of the wrong shape and a +// key that collides with unrelated configs. Deriving happens in exactly one place +// (nvte_get_fused_attn_backend_v2), so this is what keeps a new path into the builders from +// quietly skipping it. +inline void check_derived(const FusedAttnConfig &cfg) { + NVTE_CHECK( + cfg.is_derived, + "FusedAttnConfig reached a graph build with its derived fields unset. Every config " + "must pass through FusedAttnConfig::derive() first; see nvte_get_fused_attn_backend_v2."); +} + 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/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index b3b9922abf..5ef661ffe0 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -247,7 +247,16 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi const char **message) { using namespace transformer_engine; using namespace transformer_engine::fused_attn; - const FusedAttnConfig &cfg = *get_fused_attn_config(config); + // Every config entering this library passes through here on its way to a graph, so this is the + // one place that has to fill the derived fields, and it does so in place. The caller keeps + // ownership; what it gets back is its own config with the blanks filled in, which is what the + // execution path then hands to the backend it selected -- nvte_fused_attn_fwd_v2() queries with + // the very config it goes on to run, so deriving here is what lets the run reuse the graph the + // query built rather than key a second one. Deriving is idempotent, so a config that arrives + // already derived is unharmed. It is a write, though, so one config object must not be queried + // from two threads at once; every caller here builds its config as a local, one per call. + FusedAttnConfig &cfg = *get_fused_attn_config_mutable(config); + cfg.derive(); set_message(message, ""); cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); @@ -300,12 +309,14 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi std::to_string(static_cast(qkv_format)) + "."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - 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 (cfg.check_for_forward_support) { + 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 (cfg.is_training && !cfg.is_forward) { + if (cfg.is_training && cfg.check_for_backward_support) { std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); if (!bwd_reason.empty()) { set_message(message, std::move(bwd_reason)); @@ -325,12 +336,14 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - std::string fwd_reason = is_supported_f16_fwd(cfg, handle); - if (!fwd_reason.empty()) { - set_message(message, std::move(fwd_reason)); - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (cfg.check_for_forward_support) { + std::string fwd_reason = is_supported_f16_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } } - if (cfg.is_training && !cfg.is_forward) { + if (cfg.is_training && cfg.check_for_backward_support) { std::string bwd_reason = is_supported_f16_bwd(cfg, handle); if (!bwd_reason.empty()) { set_message(message, std::move(bwd_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 b7c7a349af..b07ec848ca 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 @@ -9,8 +9,6 @@ #include #include -#include -#include #include #include "../common.h" @@ -18,84 +16,81 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" +#include "graph_cache.h" #include "graph_cache_debug.h" #include "utils.h" namespace transformer_engine { namespace fused_attn { -void fused_attn_arbitrary_seqlen_fwd_impl( - 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, 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); - 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); - 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); - 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; - float scaling_factor = cfg.attn_scale; - const float dropout_probability = cfg.dropout; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_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; - - 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) || - (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_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - bool is_dropout = (is_training && dropout_probability != 0.0f); - bool is_ragged_q = cfg.is_ragged_q; - bool is_ragged_kv = cfg.is_ragged_kv; +namespace fe = cudnn_frontend; + +using SdpaF16FwdGraphAndTensors = + std::tuple, + std::shared_ptr, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // S1 + std::shared_ptr, // S2 + std::shared_ptr, // bias + std::shared_ptr, // softmax_offset + std::shared_ptr, // seq_q / cu_seq_len_q + std::shared_ptr, // seq_kv / cu_seq_len_kv + std::shared_ptr, // page_table_k + std::shared_ptr, // page_table_v + std::shared_ptr, // offset_q + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_o + std::shared_ptr, // offset_stats + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +// What the forward graph is built from beyond the config's own fields: the dimensions ragged +// layouts bucket, and the choices that depend on the cuDNN runtime version or the SM +// architecture. The build and the execution have to reach the same answer for every one of +// these -- otherwise the graph is built for different dimensions than the pointers bound to it +// describe, or with a ragged offset width the offsets are not written in -- so they are derived +// once, by derive_f16_fwd_graph_inputs, and handed to both. +struct F16FwdGraphInputs { + // Dimensions the graph is built at. Ragged layouts substitute bucketed token counts for + // max_seqlen (and, unless cu_seqlens are passed to cuDNN directly, a bucketed batch size) + // so that one graph serves every shape that falls in the same bucket. + int64_t b; + int64_t s_q; + int64_t s_kv; + // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever + // the bucketing above did to `b`. + int64_t actual_b; + bool use_ragged_stats; + DType ragged_offset_type; + RaggedOffsetMultipliers offset_mults; +}; + +// Derives the above, and rejects configurations that no graph can serve. Those rejections +// depend on combinations of fields rather than any single one, so they cannot live in the +// config's own validation; running them here is what lets a support query answer for them +// without building anything. +static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) { + check_derived(cfg); + const bool is_padding = cfg.is_padding; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; 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; + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - bool is_paged_kv = cfg.is_paged_kv; - if (is_paged_kv) { + if (cfg.is_paged_kv) { NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } - // 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 = cfg.uses_cu_seqlens_directly; - + int64_t b = static_cast(cfg.batch_size); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.max_seqlen_kv); // keep original batch size because cu_seqlens are created with [b+1] shape - int64_t actual_b = b; + const int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] @@ -108,364 +103,401 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // cuDNN reads the user's [actual_b+1] cu_seqlens buffers, so a quantized batch // would read out of bounds. if (!use_cu_seqlens_directly) { - b = bucketed_batch_size; + b = static_cast(cfg.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; + s_q = is_ragged_q ? static_cast(cfg.bucketed_num_tokens_q) : s_q; + s_kv = is_ragged_kv ? static_cast(cfg.bucketed_num_tokens_kv) : s_kv; } } + const bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; const DType ragged_offset_type = use_cu_seqlens_directly ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. : (cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32); - // Ragged offset multipliers (elements per token); shared with the legacy conversion // kernel (cu_seqlens_padded_to_offsets) so the two paths cannot drift apart. - const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + const RaggedOffsetMultipliers offset_mults( + nvte_get_qkv_layout_group(cfg.qkv_layout), static_cast(cfg.num_attn_heads), + static_cast(cfg.num_gqa_groups), static_cast(cfg.head_dim_qk), + static_cast(cfg.head_dim_v)); + + // Field order must match F16FwdGraphInputs; one per line so that it can be checked by eye. + return F16FwdGraphInputs{ + b, s_q, s_kv, actual_b, use_ragged_stats, ragged_offset_type, offset_mults, + }; +} - bool generate_stats = true; // Always return stats - const FusedAttnConfig cache_cfg = cfg.make_cache_key(); - try { - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // S1 - std::shared_ptr, // S2 - std::shared_ptr, // bias - std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q / cu_seq_len_q - std::shared_ptr, // seq_kv / cu_seq_len_kv - std::shared_ptr, // page_table_k - std::shared_ptr, // page_table_v - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - // 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 { - // Lock the map lookup, not the build, so different graphs can build in parallel - 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; - } - graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); - if (cache_hit) { - return cached_graph; - } +// Constructs the forward graph for one cache key, and only constructs it: whether cuDNN will run +// it is settled by the caller, in get_or_build_cached_graph(), which is also where the plan build +// eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every call +// that does need one now sits on the other side of that boundary. +// +// Everything the graph's shape and topology depends on comes from `cfg` and `in`, so the build +// has one source of truth and cannot drift from the caller that will bind pointers to it. +static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig &cfg, + const F16FwdGraphInputs &in) { + const int64_t b = in.b; + const int64_t s_q = in.s_q; + const int64_t s_kv = in.s_kv; + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + const int64_t num_pages_k = static_cast(cfg.num_pages_k); + const int64_t num_pages_v = static_cast(cfg.num_pages_v); + const int64_t page_size_k = static_cast(cfg.page_size_k); + const int64_t page_size_v = static_cast(cfg.page_size_v); + const int64_t max_pages_per_seq_k = static_cast(cfg.max_pages_per_seq_k); + const int64_t max_pages_per_seq_v = static_cast(cfg.max_pages_per_seq_v); + const int64_t bias_b = static_cast(cfg.bias_batch_size); + const int64_t bias_h = static_cast(cfg.bias_num_heads); + const int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool is_training = cfg.is_training; + const bool return_max_logit = cfg.return_max_logit; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); + const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_paged_kv = cfg.is_paged_kv; + const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_dropout = (is_training && dropout_probability != 0.0f); + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; + const auto cudnn_runtime_version = cudnnGetVersion(); + const bool use_ragged_stats = in.use_ragged_stats; + const DType ragged_offset_type = in.ragged_offset_type; + const RaggedOffsetMultipliers offset_mults = in.offset_mults; + const bool generate_stats = true; // Always return stats + + auto mha_graph = std::make_shared(); + mha_graph->set_io_data_type(tensorType) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr Q, K, V, attn_scale, softmax_offset; + std::shared_ptr bias, seq_q, seq_kv; + std::shared_ptr page_table_k, page_table_v; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; + std::shared_ptr dropout_seed, dropout_offset; + + std::vector q_stride(4); + std::vector k_stride(4); + std::vector v_stride(4); + generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_Q_Matrix); + if (is_paged_kv) { + generateMatrixStrides(num_pages_k, hg, page_size_k, page_size_v, d_qk, k_stride.data(), + qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); + generateMatrixStrides(num_pages_v, hg, page_size_k, page_size_v, d_v, v_stride.data(), + qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); + } else { + generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_K_Matrix); + generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_V_Matrix); + } - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(tensorType) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, K, V, attn_scale, softmax_offset; - std::shared_ptr bias, seq_q, seq_kv; - std::shared_ptr page_table_k, page_table_v; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - if (is_paged_kv) { - generateMatrixStrides(num_pages_k, hg, page_size_k, page_size_v, d_qk, k_stride.data(), - qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(num_pages_v, hg, page_size_k, page_size_v, d_v, v_stride.data(), - qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - } else { - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - } + Q = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d_qk}).set_stride(q_stride)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + if (use_cu_seqlens_directly) { + Q->set_ragged_offset_multiplier(offset_mults.q); + } + } + K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); + V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); + if (is_paged_kv) { + K->set_dim({num_pages_k, hg, page_size_k, d_qk}); + V->set_dim({num_pages_v, hg, page_size_v, d_v}); + } else if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); + V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); + if (use_cu_seqlens_directly) { + K->set_ragged_offset_multiplier(offset_mults.k); + V->set_ragged_offset_multiplier(offset_mults.v); + } + } else { + K->set_dim({b, hg, s_kv, d_qk}); + V->set_dim({b, hg, s_kv, d_v}); + } - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_stride)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - if (use_cu_seqlens_directly) { - Q->set_ragged_offset_multiplier(offset_mults.q); - } - } - K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); - V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); - if (is_paged_kv) { - K->set_dim({num_pages_k, hg, page_size_k, d_qk}); - V->set_dim({num_pages_v, hg, page_size_v, d_v}); - } else if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); - V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); - if (use_cu_seqlens_directly) { - K->set_ragged_offset_multiplier(offset_mults.k); - V->set_ragged_offset_multiplier(offset_mults.v); - } - } else { - K->set_dim({b, hg, s_kv, d_qk}); - V->set_dim({b, hg, s_kv, d_v}); - } + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); + + fe::graph::SDPA_attributes sdpa_options; + sdpa_options = fe::graph::SDPA_attributes() + .set_name("flash_attention") + .set_generate_stats(generate_stats) + .set_attn_scale(attn_scale); + + fe::DiagonalAlignment_t const &diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_options.set_diagonal_alignment(diagonal_alignment); + if (cudnn_runtime_version >= 90200 && window_size_left != -1) { + sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (cudnn_runtime_version >= 90600 && window_size_right != -1) { + sdpa_options.set_diagonal_band_right_bound(window_size_right); + } + if (is_causal || is_causal_bottom_right) { + sdpa_options.set_diagonal_band_right_bound(0); + } + + sdpa_options.set_alibi_mask(is_alibi); + + if (is_bias) { + bias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("bias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + sdpa_options.set_bias(bias); + } + + if (is_padding) { + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_cu_seq_len_q(seq_q).set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + } + + if (is_paged_kv) { + page_table_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("page_table_k") + .set_dim({b, 1, max_pages_per_seq_k, 1}) + .set_stride({{max_pages_per_seq_k, max_pages_per_seq_v, 1, 1}}) + .set_data_type(fe::DataType_t::INT32)); + page_table_v = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("page_table_v") + .set_dim({b, 1, max_pages_per_seq_v, 1}) + .set_stride({{max_pages_per_seq_v, max_pages_per_seq_v, 1, 1}}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_paged_attention_k_table(page_table_k); + sdpa_options.set_paged_attention_v_table(page_table_v); + sdpa_options.set_paged_attention_max_seq_len_kv(static_cast(s_kv)); + } - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - fe::graph::SDPA_attributes sdpa_options; - sdpa_options = fe::graph::SDPA_attributes() - .set_name("flash_attention") - .set_generate_stats(generate_stats) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const &diagonal_alignment = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_options.set_diagonal_alignment(diagonal_alignment); - if (cudnn_runtime_version >= 90200 && window_size_left != -1) { - sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (cudnn_runtime_version >= 90600 && window_size_right != -1) { - sdpa_options.set_diagonal_band_right_bound(window_size_right); - } - if (is_causal || is_causal_bottom_right) { - sdpa_options.set_diagonal_band_right_bound(0); - } - - sdpa_options.set_alibi_mask(is_alibi); + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - if (is_bias) { - bias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_options.set_bias(bias); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_options.set_sink_token(softmax_offset); + } - if (is_padding) { - if (use_cu_seqlens_directly) { - // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_kv") + std::shared_ptr Max; + if (use_ragged_stats) { + offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") .set_dim({b + 1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding) - .set_cu_seq_len_q(seq_q) - .set_cu_seq_len_kv(seq_kv); - // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. - // Pin the implementation so an unsupported config fails with the unified - // engine's specific error instead of auto-selection's generic failure. - sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); - } else { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - } - - if (is_paged_kv) { - page_table_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("page_table_k") - .set_dim({b, 1, max_pages_per_seq_k, 1}) - .set_stride({{max_pages_per_seq_k, max_pages_per_seq_v, 1, 1}}) - .set_data_type(fe::DataType_t::INT32)); - page_table_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("page_table_v") - .set_dim({b, 1, max_pages_per_seq_v, 1}) - .set_stride({{max_pages_per_seq_v, max_pages_per_seq_v, 1, 1}}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_paged_attention_k_table(page_table_k); - sdpa_options.set_paged_attention_v_table(page_table_v); - sdpa_options.set_paged_attention_max_seq_len_kv(static_cast(s_kv)); + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } + if (return_max_logit) { + Max = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Max") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Max->set_ragged_offset_multiplier(offset_mults.stats); } + } else { + Max->set_stride({h * s_q, s_q, 1, 1}); + } + sdpa_options.set_logit_max(Max); + } - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } + auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); + + std::vector o_stride(4); + generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_O_Matrix); + O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_stride); + if (is_ragged_q) { + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + O->set_ragged_offset(offset_o); + if (use_cu_seqlens_directly) { + O->set_ragged_offset_multiplier(offset_mults.o); + } + } - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_options.set_sink_token(softmax_offset); - } + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Stats->set_ragged_offset_multiplier(offset_mults.stats); + } + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } - std::shared_ptr Max; - if (use_ragged_stats) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - if (return_max_logit) { - Max = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Max") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Max->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Max->set_stride({h * s_q, s_q, 1, 1}); - } - sdpa_options.set_logit_max(Max); - } + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // attn_scale + std::shared_ptr> // O + key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); + auto Stats_tuple = + return_max_logit ? std::make_tuple(Stats, Max) : std::make_tuple(Stats, nullptr); + auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); + auto softmax_offset_tuple = + is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto page_table_tuple = + is_paged_kv ? std::make_tuple(page_table_k, page_table_v) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return 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); +} - auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); +// The forward graph cache and the only route to it. Both the execution path and the support +// probe come through here, so a probe leaves behind exactly the entry a later execution finds. +// That is what lets the probe's answer describe the graph that actually runs, rather than a +// separately built lookalike. +static std::shared_ptr> f16_fwd_cached_graph( + const FusedAttnConfig &cfg, const F16FwdGraphInputs &in, cudnnHandle_t handle) { + static GraphCache cache; + return get_or_build_cached_graph(cache, cfg.make_cache_key(), "fwd", handle, + [&] { return build_sdpa_f16_fwd_graph(cfg, in); }); +} - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_stride); - if (is_ragged_q) { - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - O->set_ragged_offset(offset_o); - if (use_cu_seqlens_directly) { - O->set_ragged_offset_multiplier(offset_mults.o); - } - } +void fused_attn_arbitrary_seqlen_fwd_impl( + 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, void *workspace, size_t *workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { + using namespace transformer_engine; - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Stats->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } + // Derived once and handed to the cache, which passes them to the graph build, so that the + // graph and the pointers bound to it below cannot be decided differently. Also where an + // unserviceable configuration is rejected. + const F16FwdGraphInputs in = derive_f16_fwd_graph_inputs(cfg); + const int64_t b = in.b; + const int64_t actual_b = in.actual_b; + const bool use_ragged_stats = in.use_ragged_stats; + const DType ragged_offset_type = in.ragged_offset_type; + const RaggedOffsetMultipliers offset_mults = in.offset_mults; - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // attn_scale - std::shared_ptr> // O - key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); - auto Stats_tuple = - return_max_logit ? std::make_tuple(Stats, Max) : std::make_tuple(Stats, nullptr); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto page_table_tuple = is_paged_kv ? std::make_tuple(page_table_k, page_table_v) - : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = - use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : 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})); - }); - 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. - { - std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); - auto inserted = cache.insert({descriptor, return_tuple}); - return inserted.first->second; - } - }; + const bool return_max_logit = cfg.return_max_logit; + // Not const: bound into the variant pack by address as a pass-by-value graph input. + float scaling_factor = cfg.attn_scale; + const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_padding = cfg.is_padding; + const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool is_paged_kv = cfg.is_paged_kv; + // 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 = cfg.uses_cu_seqlens_directly; + try { + auto cache_entry = f16_fwd_cached_graph(cfg, in, handle); 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, cache_cfg); + dropout_seed, dropout_offset] = cache_entry->tensors; + + // This graph is going to be used, so finish the build the cache deferred. + ensure_plans_built("fwd", *cache_entry); // 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. @@ -607,391 +639,410 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } -} // NOLINT(readability/fn_size) - -void fused_attn_arbitrary_seqlen_bwd_impl( - 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)); +using SdpaF16BwdGraphAndTensors = + std::tuple, + std::shared_ptr, // q + std::shared_ptr, // k + std::shared_ptr, // v + std::shared_ptr, // o + std::shared_ptr, // dO + std::shared_ptr, // stats + std::shared_ptr, // attn_scale + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr, // dV + std::shared_ptr, // bias + std::shared_ptr, // dBias + std::shared_ptr, // softmax_offset + std::shared_ptr, // d_softmax_offset + std::shared_ptr, // seq_q + std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_o + std::shared_ptr, // offset_stats + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +// The backward equivalent of F16FwdGraphInputs; see there for why these are derived once and +// shared. The backward graph reads no page table and passes no cu_seqlens straight through, so +// it needs neither the paged-attention check nor the ragged offset multipliers. +struct F16BwdGraphInputs { + int64_t b; + int64_t s_q; + int64_t s_kv; + int64_t actual_b; + bool use_ragged_stats; + DType ragged_offset_type; +}; + +// The backward counterpart of derive_f16_fwd_graph_inputs; see there for what the rejections are +// doing here and why a support query can answer for them without building a graph. +static F16BwdGraphInputs derive_f16_bwd_graph_inputs(const FusedAttnConfig &cfg) { + check_derived(cfg); + const bool is_padding = cfg.is_padding; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); 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); int64_t s_q = static_cast(cfg.max_seqlen_q); 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); - 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); - float scaling_factor = cfg.attn_scale; - const float dropout_probability = cfg.dropout; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_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) || - (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_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - bool is_dropout = (dropout_probability != 0.0f); - 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; - // keep original batch size because cu_seqlens are created with [b+1] shape - int64_t actual_b = b; + const int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). 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 = 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; + // for query and key/value so the graph is static within each quantization bucket. + // The batch is bucketed unconditionally here, where the forward pass guards it: only + // the forward graph can be handed the user's cu_seqlens buffers directly, and it is + // their [actual_b+1] length that a quantized batch would overrun. The backward graph + // always reads converted seqlens out of our own workspace, so nothing here is sized by + // the true batch. make_cache_key() splits on the pass for this reason as well. + b = static_cast(cfg.bucketed_batch_size); + s_q = is_ragged_q ? static_cast(cfg.bucketed_num_tokens_q) : s_q; + s_kv = is_ragged_kv ? static_cast(cfg.bucketed_num_tokens_kv) : s_kv; } } + + const bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; // 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(); - try { - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // dO - std::shared_ptr, // stats - std::shared_ptr, // attn_scale - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // bias - std::shared_ptr, // dBias - std::shared_ptr, // softmax_offset - std::shared_ptr, // d_softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - 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 { - // Lock the map lookup, not the build, so different graphs can build in parallel - 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; - } - graph_cache_debug::record_cache_lookup("bwd", cache_hit, cfg); - if (cache_hit) { - return cached_graph; - } + // Field order must match F16BwdGraphInputs; one per line so that it can be checked by eye. + return F16BwdGraphInputs{ + b, s_q, s_kv, actual_b, use_ragged_stats, ragged_offset_type, + }; +} - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(tensorType) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr q, k, v, o, dO, stats, attn_scale; - std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset, - seq_q, seq_kv; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - - q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_stride)); - k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_stride)); - v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_stride)); - o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_stride)); - dO = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_stride)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - q->set_ragged_offset(offset_q); - o->set_ragged_offset(offset_o); - dO->set_ragged_offset(offset_o); - } - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") +// The backward counterpart of build_sdpa_f16_fwd_graph; see there for why it constructs the graph +// and nothing else. +// +// Everything the graph's shape and topology depends on comes from `cfg` and `in`, so the build +// has one source of truth and cannot drift from the caller that will bind pointers to it. +static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig &cfg, + const F16BwdGraphInputs &in) { + const int64_t b = in.b; + const int64_t s_q = in.s_q; + const int64_t s_kv = in.s_kv; + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + const int64_t bias_b = static_cast(cfg.bias_batch_size); + const int64_t bias_h = static_cast(cfg.bias_num_heads); + const int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); + const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_dropout = (dropout_probability != 0.0f); + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); + const bool use_ragged_stats = in.use_ragged_stats; + const DType ragged_offset_type = in.ragged_offset_type; + + auto mha_graph = std::make_shared(); + mha_graph->set_io_data_type(tensorType) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr q, k, v, o, dO, stats, attn_scale; + std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset, + seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; + std::shared_ptr dropout_seed, dropout_offset; + + std::vector q_stride(4); + std::vector k_stride(4); + std::vector v_stride(4); + std::vector o_stride(4); + generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_Q_Matrix); + generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_K_Matrix); + generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_V_Matrix); + generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_O_Matrix); + + q = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d_qk}).set_stride(q_stride)); + k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_stride)); + v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_stride)); + o = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("O").set_dim({b, h, s_q, d_v}).set_stride(o_stride)); + dO = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("dO").set_dim({b, h, s_q, d_v}).set_stride(o_stride)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + q->set_ragged_offset(offset_q); + o->set_ragged_offset(offset_o); + dO->set_ragged_offset(offset_o); + } + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + k->set_ragged_offset(offset_k); + v->set_ragged_offset(offset_v); + } + + stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("stats") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") .set_dim({b + 1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - k->set_ragged_offset(offset_k); - v->set_ragged_offset(offset_v); - } + stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + stats->set_stride({h * s_q, s_q, 1, 1}); + } - stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("stats") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - stats->set_stride({h * s_q, s_q, 1, 1}); - } + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); + fe::graph::SDPA_backward_attributes sdpa_backward_options; + sdpa_backward_options = fe::graph::SDPA_backward_attributes() + .set_name("flash_attention_backward") + .set_attn_scale(attn_scale); - fe::graph::SDPA_backward_attributes sdpa_backward_options; - sdpa_backward_options = fe::graph::SDPA_backward_attributes() - .set_name("flash_attention_backward") - .set_attn_scale(attn_scale); + if (use_ragged_stats) { + sdpa_backward_options.set_max_total_seq_len_q(s_q); + } + if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + sdpa_backward_options.set_max_total_seq_len_kv(s_kv); + } - if (use_ragged_stats) { - sdpa_backward_options.set_max_total_seq_len_q(s_q); - } - if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { - sdpa_backward_options.set_max_total_seq_len_kv(s_kv); - } + fe::DiagonalAlignment_t const &diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); - fe::DiagonalAlignment_t const &diagonal_alignment = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); + if (cudnn_runtime_version >= 90200 && window_size_left != -1) { + sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (cudnn_runtime_version >= 90600 && window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } + if (is_causal || is_causal_bottom_right) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } - if (cudnn_runtime_version >= 90200 && window_size_left != -1) { - sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (cudnn_runtime_version >= 90600 && window_size_right != -1) { - sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } - if (is_causal || is_causal_bottom_right) { - sdpa_backward_options.set_diagonal_band_right_bound(0); - } + if (cudnn_runtime_version >= 90000) { + sdpa_backward_options.set_deterministic_algorithm(deterministic); + } - if (cudnn_runtime_version >= 90000) { - sdpa_backward_options.set_deterministic_algorithm(deterministic); - } + sdpa_backward_options.set_alibi_mask(is_alibi); + + if (is_bias) { + bias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("bias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + sdpa_backward_options.set_bias(bias); + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + dBias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("dBias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + sdpa_backward_options.set_dbias(dBias); + } + } - sdpa_backward_options.set_alibi_mask(is_alibi); - - if (is_bias) { - bias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_backward_options.set_bias(bias); - // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation - // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 - if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { - dBias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("dBias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_backward_options.set_dbias(dBias); - } - } + if (is_padding) { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } - if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_backward_options.set_padding_mask(is_padding) - .set_seq_len_q(seq_q) - .set_seq_len_kv(seq_kv); - } + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_sink_token(softmax_offset); + d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("d_softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_dsink_token(d_softmax_offset); + } - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_sink_token(softmax_offset); - d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("d_softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_dsink_token(d_softmax_offset); - } + auto [dQ, dK, dV] = mha_graph->sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options); - auto [dQ, dK, dV] = mha_graph->sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options); + dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(q_stride); + dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(k_stride); + dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(v_stride); + if (is_ragged_q) { + dQ->set_ragged_offset(offset_q); + } + if (is_ragged_kv) { + dK->set_ragged_offset(offset_k); + dV->set_ragged_offset(offset_v); + } - dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(q_stride); - dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(k_stride); - dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(v_stride); - if (is_ragged_q) { - dQ->set_ragged_offset(offset_q); - } - if (is_ragged_kv) { - dK->set_ragged_offset(offset_k); - dV->set_ragged_offset(offset_v); - } + std::tuple, // q + std::shared_ptr, // k + std::shared_ptr, // v + std::shared_ptr, // o + std::shared_ptr, // dO + std::shared_ptr, // stats + std::shared_ptr, // attn_scale + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr> // dV + key_tensors_tuple = std::make_tuple(q, k, v, o, dO, stats, attn_scale, dQ, dK, dV); + auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); + auto softmax_offset_tuple = is_softmax_offset ? std::make_tuple(softmax_offset, d_softmax_offset) + : std::make_tuple(nullptr, nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return 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); +} - std::tuple, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // dO - std::shared_ptr, // stats - std::shared_ptr, // attn_scale - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr> // dV - key_tensors_tuple = std::make_tuple(q, k, v, o, dO, stats, attn_scale, dQ, dK, dV); - auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); - auto softmax_offset_tuple = is_softmax_offset - ? std::make_tuple(softmax_offset, d_softmax_offset) - : std::make_tuple(nullptr, nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = - use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : 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})); - }); - 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. - { - std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); - auto inserted = cache.insert({descriptor, return_tuple}); - return inserted.first->second; - } - }; +// The backward counterpart of f16_fwd_cached_graph; see there. +static std::shared_ptr> f16_bwd_cached_graph( + const FusedAttnConfig &cfg, const F16BwdGraphInputs &in, cudnnHandle_t handle) { + static GraphCache cache; + return get_or_build_cached_graph(cache, cfg.make_cache_key(), "bwd", handle, + [&] { return build_sdpa_f16_bwd_graph(cfg, in); }); +} + +void fused_attn_arbitrary_seqlen_bwd_impl( + 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; + + // Derived once and handed to the cache, which passes them to the graph build, so that the + // graph and the pointers bound to it below cannot be decided differently. Also where an + // unserviceable configuration is rejected. + const F16BwdGraphInputs in = derive_f16_bwd_graph_inputs(cfg); + const int64_t b = in.b; + const int64_t actual_b = in.actual_b; + const bool use_ragged_stats = in.use_ragged_stats; + const DType ragged_offset_type = in.ragged_offset_type; + + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + // Not const: bound into the variant pack by address as a pass-by-value graph input. + float scaling_factor = cfg.attn_scale; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_padding = cfg.is_padding; + const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_dropout = (cfg.dropout != 0.0f); + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + try { + auto cache_entry = f16_bwd_cached_graph(cfg, in, handle); 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, cache_cfg); + dropout_seed, dropout_offset] = cache_entry->tensors; + + // This graph is going to be used, so finish the build the cache deferred. + ensure_plans_built("bwd", *cache_entry); // 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. @@ -1165,9 +1216,6 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i void *devPtrPageTableK = page_table_k ? page_table_k->data.dptr : nullptr; void *devPtrPageTableV = page_table_v ? page_table_v->data.dptr : nullptr; - FusedAttnConfig graph_cfg = cfg; - graph_cfg.derive(); - size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); @@ -1204,8 +1252,8 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i 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 = {graph_cfg.bias_batch_size, graph_cfg.bias_num_heads, - graph_cfg.bias_seqlen_q, graph_cfg.bias_seqlen_kv}; + output_bias->data.shape = {cfg.bias_batch_size, cfg.bias_num_heads, cfg.bias_seqlen_q, + cfg.bias_seqlen_kv}; output_bias->data.dtype = QKV_type; } @@ -1246,10 +1294,10 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i size_t workspace_size = 0; fused_attn_arbitrary_seqlen_fwd_impl( - 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); + 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) { if (workspace->data.dptr == nullptr) { @@ -1293,9 +1341,6 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i devPtrdBias = output_dBias->data.dptr; } - FusedAttnConfig graph_cfg = cfg; - graph_cfg.derive(); - void *devPtrdQ = output_dQ->data.dptr; void *devPtrdK = output_dK->data.dptr; void *devPtrdV = output_dV->data.dptr; @@ -1320,11 +1365,10 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i size_t workspace_size = 0; fused_attn_arbitrary_seqlen_bwd_impl( - 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); + 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) { @@ -1341,53 +1385,49 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i } } +// Whether cuDNN can run the forward graph this config asks for: the empty string if it can, +// otherwise cuDNN's own account of why not, which the backend selector reports to the caller. +// +// The question is answered by deriving the graph's inputs and building the graph, which is +// where every rejection comes from -- there is no separate list of rules to keep in step with +// the builder. The graph goes into the same cache the execution path reads, so the work is not +// thrown away and what was checked is what will run. It stops short of build_plans(), the +// expensive step, which the first execution of the graph does instead; see CachedGraph. +// +// A refusal is cached too, so asking the same question twice costs one build rather than two; +// the second answer is the first one replayed. See GraphCache. +// +// The copy below is made for the sake of one flag, which is not a redundant restatement of what +// the caller already asked for: make_cache_key() reads it to choose between the forward and the +// backward normalization, and one config can be probed in both directions -- the deprecated +// nvte_get_fused_attn_backend() leaves both check_for_*_support set, so both probes run off a +// single config. Each probe therefore states its own direction instead of inheriting it. std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; - graph_cfg.is_forward = true; - graph_cfg.derive(); + graph_cfg.check_for_forward_support = true; - size_t workspace_size = 0; try { - fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( - graph_cfg, - /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrBias=*/nullptr, - /*devPtrSoftmaxOffset=*/nullptr, /*devPtrS1=*/nullptr, /*devPtrS2=*/nullptr, - /*devPtrO=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, - /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, - /*devPtrPageTableK=*/nullptr, /*devPtrPageTableV=*/nullptr, - /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, - /*workspace=*/nullptr, &workspace_size, - /*stream=*/static_cast(0), handle); + const fused_attn::F16FwdGraphInputs in = fused_attn::derive_f16_fwd_graph_inputs(graph_cfg); + fused_attn::f16_fwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception &e) { - return e.what(); + return fused_attn::refusal_reason(e, "is_supported_f16_fwd: rejected without a reason."); } catch (...) { return "is_supported_f16_fwd: unknown failure."; } } +// The backward counterpart of is_supported_f16_fwd; see there. std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; - graph_cfg.is_forward = false; - graph_cfg.derive(); + graph_cfg.check_for_forward_support = false; - size_t workspace_size = 0; try { - fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( - graph_cfg, - /*devPtrQ=*/nullptr, /*devPtrKTranspose=*/nullptr, - /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, - /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, - /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, /*devPtrdO=*/nullptr, - /*devPtrdBias=*/nullptr, /*devPtrdSoftmaxOffset=*/nullptr, - /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, - /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, - /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, - /*workspace=*/nullptr, &workspace_size, - /*stream=*/static_cast(0), handle); + const fused_attn::F16BwdGraphInputs in = fused_attn::derive_f16_bwd_graph_inputs(graph_cfg); + fused_attn::f16_bwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception &e) { - return e.what(); + return fused_attn::refusal_reason(e, "is_supported_f16_bwd: rejected without a reason."); } catch (...) { return "is_supported_f16_bwd: unknown failure."; } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 9449f74206..ad18a19481 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -4,13 +4,13 @@ * See LICENSE for license information. ************************************************************************/ -#include #include #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" +#include "graph_cache.h" #include "graph_cache_debug.h" #include "utils.h" @@ -18,64 +18,67 @@ namespace transformer_engine { namespace fused_attn { using namespace transformer_engine; +namespace fe = cudnn_frontend; // 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) { - using namespace transformer_engine; +using SdpaFp8FwdGraphAndTensors = + std::tuple, + std::shared_ptr, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_s + std::shared_ptr, // scale_s + std::shared_ptr, // scale_o + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // amax_s + std::shared_ptr, // amax_o + std::shared_ptr, // Stats + std::shared_ptr, // bias + std::shared_ptr, // softmax_offset + std::shared_ptr, // seq_q + std::shared_ptr, // seq_kv + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +// The FP8 forward path's policy decisions: which quantization recipe the graph is built for, +// and whether cu_seqlens can be handed to cuDNN directly. Both decide which tensors the graph +// has, and so which pointers the variant pack has to bind -- the build and the execution cannot +// answer them differently, which is why they are derived once, here, for both. +struct Fp8FwdGraphInputs { + bool is_delayed_scaling; + bool is_current_scaling; + bool is_mxfp8; + bool use_cu_seqlens_directly; +}; + +// Derives the above and rejects what FP8 cannot serve. Unlike the F16 path there is no +// bucketing to do, because FP8 has no ragged/THD support: the graph's shapes are exactly the +// config's. +static Fp8FwdGraphInputs derive_fp8_fwd_graph_inputs(const FusedAttnConfig& cfg) { + check_derived(cfg); 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 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; - 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) || - (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); - 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) && - (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_alibi = (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI); + const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); + + 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!"); + const bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || + o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); + const bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + const bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); NVTE_CHECK( is_delayed_scaling || is_current_scaling || is_mxfp8, "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); @@ -100,329 +103,344 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de // (which doesn't support cu_seqlens). Remove this restriction when possible. !is_dropout; - const FusedAttnConfig cache_cfg = cfg.make_cache_key(); - try { - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr, // amax_o - std::shared_ptr, // Stats - std::shared_ptr, // bias - std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - // 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 { - // Lock the map lookup, not the build, so different graphs can build in parallel - 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; - } - graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); - if (cache_hit) { - return cached_graph; - } + // Field order must match Fp8FwdGraphInputs; one per line so that it can be checked by eye. + return Fp8FwdGraphInputs{ + is_delayed_scaling, + is_current_scaling, + is_mxfp8, + use_cu_seqlens_directly, + }; +} - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(qkv_tensor_type) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, K, V, attn_scale; - std::shared_ptr descale_q, descale_k, descale_v; - std::shared_ptr descale_s, scale_s, scale_o; - std::shared_ptr bias, softmax_offset, seq_q, seq_kv; - std::shared_ptr dropout_seed, dropout_offset; - - // Q, K, V, attn_scale - std::vector q_strides(4), k_strides(4), v_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), - k_strides.data(), v_strides.data(), qkv_layout); - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_strides) - .set_data_type(qkv_tensor_type)); - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) - .set_data_type(qkv_tensor_type)); - V = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_strides) - .set_data_type(qkv_tensor_type)); - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") +// Constructs the forward FP8 graph for one cache key, and only constructs it: whether cuDNN will +// run it is settled by the caller, in get_or_build_cached_graph(), which is also where the plan +// build eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every +// call that does need one now sits on the other side of that boundary. +// +// Everything the graph's shape and topology depends on comes from `cfg` and `in`, so the build +// has one source of truth and cannot drift from the caller that will bind pointers to it. +static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& cfg, + const Fp8FwdGraphInputs& in) { + 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 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 int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool is_training = cfg.is_training; + 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 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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_dropout = (is_training && dropout_probability != 0.0f); + const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_delayed_scaling = in.is_delayed_scaling; + const bool is_current_scaling = in.is_current_scaling; + const bool is_mxfp8 = in.is_mxfp8; + const bool use_cu_seqlens_directly = in.use_cu_seqlens_directly; + + auto mha_graph = std::make_shared(); + mha_graph->set_io_data_type(qkv_tensor_type) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr Q, K, V, attn_scale; + std::shared_ptr descale_q, descale_k, descale_v; + std::shared_ptr descale_s, scale_s, scale_o; + std::shared_ptr bias, softmax_offset, seq_q, seq_kv; + std::shared_ptr dropout_seed, dropout_offset; + + // Q, K, V, attn_scale + std::vector q_strides(4), k_strides(4), v_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), + k_strides.data(), v_strides.data(), qkv_layout); + Q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q") + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_strides) + .set_data_type(qkv_tensor_type)); + K = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_strides) + .set_data_type(qkv_tensor_type)); + V = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_strides) + .set_data_type(qkv_tensor_type)); + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); + + // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o + if (is_delayed_scaling || is_current_scaling) { + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); + descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); + descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); + scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); + if (is_delayed_scaling) { + scale_o = mha_graph->tensor_like(descale_q, "Scale_o"); + } + if (is_current_scaling) { + 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 kv_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.kv_format; + std::vector q_scale_strides(4); + std::vector k_scale_strides(4); + std::vector v_scale_strides(4); + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_v_padded, + v_scale_strides.data(), kv_scale_inv_format); + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_v_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + } + + fe::graph::SDPA_fp8_attributes sdpa_options; + 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 = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_options.set_diagonal_band_right_bound(window_size_right); + } + } + if (is_causal_bottom_right) { + sdpa_options.set_diagonal_band_right_bound(0); + } + + // sdpa_options.set_alibi_mask(is_alibi); + // if (is_bias) { + // bias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("bias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // sdpa_options.set_bias(bias); + // } + + if (is_padding) { + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_cu_seq_len_q(seq_q).set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + } + + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o - if (is_delayed_scaling || is_current_scaling) { - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); - if (is_delayed_scaling) { - scale_o = mha_graph->tensor_like(descale_q, "Scale_o"); - } - if (is_current_scaling) { - 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 kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) - ? qkv_scale_inv_format - : cfg.kv_format; - std::vector q_scale_strides(4); - std::vector k_scale_strides(4); - std::vector v_scale_strides(4); - auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, - q_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, - k_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_v_padded, - v_scale_strides.data(), kv_scale_inv_format); - descale_q = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) - .set_stride(q_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k") - .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) - .set_stride(k_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_v") - .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_v_padded}) - .set_stride(v_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - } + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - fe::graph::SDPA_fp8_attributes sdpa_options; - 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 = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 92100) { - if (window_size_left != -1) { - sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (window_size_right != -1) { - sdpa_options.set_diagonal_band_right_bound(window_size_right); - } - } - if (is_causal_bottom_right) { - sdpa_options.set_diagonal_band_right_bound(0); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_options.set_sink_token(softmax_offset); + } - // sdpa_options.set_alibi_mask(is_alibi); - // if (is_bias) { - // bias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("bias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // sdpa_options.set_bias(bias); - // } - - if (is_padding) { - if (use_cu_seqlens_directly) { - // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_kv") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding) - .set_cu_seq_len_q(seq_q) - .set_cu_seq_len_kv(seq_kv); - // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. - // Pin the implementation so an unsupported config fails with the unified - // engine's specific error instead of auto-selection's generic failure. - sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); - } else { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - } + std::shared_ptr O, Stats, amax_s, amax_o; + if (is_delayed_scaling || is_current_scaling) { + auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, + scale_o, sdpa_options); + O = outputs[0]; + Stats = outputs[1]; + amax_s = outputs[2]; + amax_o = outputs[3]; + amax_s->set_output(true) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + } else if (is_mxfp8) { + auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, sdpa_options); + O = outputs[0]; + Stats = outputs[1]; + amax_o = outputs[2]; + } - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } + std::vector o_strides(4); + generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); + O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_strides).set_data_type(o_tensor_type); + amax_o->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + + Stats->set_output(true) + .set_data_type(fe::DataType_t::FLOAT) + .set_dim({b, h, s_q, 1}) + .set_stride({h * s_q, s_q, 1, 1}); + + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_s + std::shared_ptr, // scale_s + std::shared_ptr, // scale_o + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // amax_s + std::shared_ptr> // amax_o + key_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, + nullptr, attn_scale, O, nullptr, amax_o) + : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, + scale_o, attn_scale, O, amax_s, amax_o); + auto Stats_tuple = std::make_tuple(Stats); + auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); + auto softmax_offset_tuple = + is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, + softmax_offset_tuple, padding_tuple, dropout_tuple); +} - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_options.set_sink_token(softmax_offset); - } +// The FP8 forward graph cache and the only route to it. Both the execution path and the support +// probe come through here, so a probe leaves behind exactly the entry a later execution finds. +static std::shared_ptr> fp8_fwd_cached_graph( + const FusedAttnConfig& cfg, const Fp8FwdGraphInputs& in, cudnnHandle_t handle) { + static GraphCache cache; + return get_or_build_cached_graph(cache, cfg.make_cache_key(), "fwd", handle, + [&] { return build_sdpa_fp8_fwd_graph(cfg, in); }); +} - std::shared_ptr O, Stats, amax_s, amax_o; - if (is_delayed_scaling || is_current_scaling) { - auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, - scale_s, scale_o, sdpa_options); - O = outputs[0]; - Stats = outputs[1]; - amax_s = outputs[2]; - amax_o = outputs[3]; - amax_s->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - } else if (is_mxfp8) { - auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, sdpa_options); - O = outputs[0]; - Stats = outputs[1]; - amax_o = outputs[2]; - } +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; - std::vector o_strides(4); - generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); - O->set_output(true) - .set_dim({b, h, s_q, d_v}) - .set_stride(o_strides) - .set_data_type(o_tensor_type); - amax_o->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - - Stats->set_output(true) - .set_data_type(fe::DataType_t::FLOAT) - .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}); - - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr> // amax_o - key_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, - nullptr, attn_scale, O, nullptr, amax_o) - : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, - scale_s, scale_o, attn_scale, O, amax_s, amax_o); - auto Stats_tuple = std::make_tuple(Stats); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : 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})); - }); - 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. - { - std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); - auto inserted = cache.insert({descriptor, return_tuple}); - return inserted.first->second; - } - }; + // Derived once and handed to the cache, which passes them to the graph build, so that the + // graph and the pointers bound to it below cannot be decided differently. Also where an + // unserviceable configuration is rejected. + const Fp8FwdGraphInputs in = derive_fp8_fwd_graph_inputs(cfg); + const bool is_delayed_scaling = in.is_delayed_scaling; + const bool is_current_scaling = in.is_current_scaling; + const bool use_cu_seqlens_directly = in.use_cu_seqlens_directly; + + const int64_t b = static_cast(cfg.batch_size); + // Not const: bound into the variant pack by address as a pass-by-value graph input. + float scaling_factor = cfg.attn_scale; + const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_padding = cfg.is_padding; + const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); + const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + try { + auto cache_entry = fp8_fwd_cached_graph(cfg, in, handle); 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, cache_cfg); + dropout_offset] = cache_entry->tensors; + + // This graph is going to be used, so finish the build the cache deferred. + ensure_plans_built("fwd", *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -501,21 +519,113 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de } // 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, - 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; +using SdpaFp8BwdGraphAndTensors = + std::tuple, + std::shared_ptr, // Q + std::shared_ptr, // Q_t + std::shared_ptr, // K + std::shared_ptr, // K_t + std::shared_ptr, // V + std::shared_ptr, // O + std::shared_ptr, // Stats + std::shared_ptr, // dO + std::shared_ptr, // dO_t + std::shared_ptr, // dO_f16 + std::shared_ptr, // attn_scale + std::shared_ptr, // descale_q + std::shared_ptr, // descale_q_t + std::shared_ptr, // descale_k + std::shared_ptr, // descale_k_t + std::shared_ptr, // descale_v + std::shared_ptr, // descale_o + std::shared_ptr, // descale_dO + std::shared_ptr, // descale_dO_t + std::shared_ptr, // descale_s + std::shared_ptr, // descale_dP + std::shared_ptr, // scale_dQ + std::shared_ptr, // scale_dK + std::shared_ptr, // scale_dV + std::shared_ptr, // scale_s + std::shared_ptr, // scale_dP + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr, // dV + std::shared_ptr, // amax_dQ + std::shared_ptr, // amax_dK + std::shared_ptr, // amax_dV + std::shared_ptr, // amax_dP + std::shared_ptr, // bias + std::shared_ptr, // dBias + std::shared_ptr, // softmax_offset + std::shared_ptr, // d_softmax_offset + std::shared_ptr, // seq_q + std::shared_ptr, // seq_kv + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +// Builds the backward FP8 graph for one cache key, up to check_support() but not build_plans(); +// see CachedGraph for why the plan build is left to whoever executes the graph. +// +// Everything the graph's shape and topology depends on is re-derived from `cfg` here, so the +// build has one source of truth for them. Unlike the F16 path, FP8 has no ragged/THD support, +// so the shapes are exactly the config's and need no bucketing from the caller. +// The backward equivalent of Fp8FwdGraphInputs. The recipe is chosen from the dQKV dtype here +// rather than O's, since backward is what writes those. is_O_in_F16 additionally selects whether +// O has to be descaled on the way in. +struct Fp8BwdGraphInputs { + bool is_delayed_scaling; + bool is_current_scaling; + bool is_mxfp8; + bool is_O_in_F16; +}; + +// The backward counterpart of derive_fp8_fwd_graph_inputs; see there for the rejections and for +// why the graph's shapes are simply the config's. +static Fp8BwdGraphInputs derive_fp8_bwd_graph_inputs(const FusedAttnConfig& cfg) { + check_derived(cfg); const auto cudnn_runtime_version = cudnnGetVersion(); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + const cudnn_frontend::DataType_t dqkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.dqkv_dtype)); + const NVTEScalingMode scaling_mode = cfg.scaling_mode; + const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_alibi = (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI); + + 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!"); + const bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || + dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); + const bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || + dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + const bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || + dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + NVTE_CHECK( + is_delayed_scaling || is_current_scaling || is_mxfp8, + "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); + NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, + "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + + const bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + + // Field order must match Fp8BwdGraphInputs; one per line so that it can be checked by eye. + return Fp8BwdGraphInputs{ + is_delayed_scaling, + is_current_scaling, + is_mxfp8, + is_O_in_F16, + }; +} +// The backward counterpart of build_sdpa_fp8_fwd_graph; see there for why it constructs the graph +// and nothing else. +static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& cfg, + const Fp8BwdGraphInputs& in) { + 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 = @@ -524,7 +634,6 @@ void fused_attn_fp8_bwd_impl( 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)); - 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); @@ -532,524 +641,455 @@ 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); - float scaling_factor = cfg.attn_scale; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; const float dropout_probability = cfg.dropout; const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_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 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) || - (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); - 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) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - NVTE_CHECK( - is_delayed_scaling || is_current_scaling || is_mxfp8, - "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); - NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, - "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - - 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(); - try { - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // Q_t - std::shared_ptr, // K - std::shared_ptr, // K_t - std::shared_ptr, // V - std::shared_ptr, // O - std::shared_ptr, // Stats - std::shared_ptr, // dO - std::shared_ptr, // dO_t - std::shared_ptr, // dO_f16 - std::shared_ptr, // attn_scale - std::shared_ptr, // descale_q - std::shared_ptr, // descale_q_t - std::shared_ptr, // descale_k - std::shared_ptr, // descale_k_t - std::shared_ptr, // descale_v - std::shared_ptr, // descale_o - std::shared_ptr, // descale_dO - std::shared_ptr, // descale_dO_t - std::shared_ptr, // descale_s - std::shared_ptr, // descale_dP - std::shared_ptr, // scale_dQ - std::shared_ptr, // scale_dK - std::shared_ptr, // scale_dV - std::shared_ptr, // scale_s - std::shared_ptr, // scale_dP - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // amax_dQ - std::shared_ptr, // amax_dK - std::shared_ptr, // amax_dV - std::shared_ptr, // amax_dP - std::shared_ptr, // bias - std::shared_ptr, // dBias - std::shared_ptr, // softmax_offset - std::shared_ptr, // d_softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - 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 { - // Lock the map lookup, not the build, so different graphs can build in parallel - 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; - } - graph_cache_debug::record_cache_lookup("bwd", cache_hit, cfg); - if (cache_hit) { - return cached_graph; - } - - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - - mha_graph->set_io_data_type(qkv_tensor_type) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, Q_t, K, K_t, V, O, dO, dO_t, dO_f16, Stats, - attn_scale; - std::shared_ptr descale_q, descale_q_t, descale_k, descale_k_t, - descale_v; - std::shared_ptr descale_s, descale_o; - std::shared_ptr descale_dP, descale_dO, descale_dO_t; - std::shared_ptr scale_s, scale_dP; - std::shared_ptr scale_dQ, scale_dK, scale_dV; - std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; - std::shared_ptr seq_q, seq_kv; - std::shared_ptr dropout_seed, dropout_offset; - - // Q, K, V, O, dO, stats, attn_scale - std::vector q_strides(4), k_strides(4), v_strides(4), o_strides(4), dO_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), - k_strides.data(), v_strides.data(), qkv_layout); - generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); - generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_strides.data(), do_format); - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") + const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_dropout = (dropout_probability != 0.0f); + const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_delayed_scaling = in.is_delayed_scaling; + const bool is_current_scaling = in.is_current_scaling; + const bool is_mxfp8 = in.is_mxfp8; + const bool is_O_in_F16 = in.is_O_in_F16; + + auto mha_graph = std::make_shared(); + + mha_graph->set_io_data_type(qkv_tensor_type) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr Q, Q_t, K, K_t, V, O, dO, dO_t, dO_f16, Stats, + attn_scale; + std::shared_ptr descale_q, descale_q_t, descale_k, descale_k_t, + descale_v; + std::shared_ptr descale_s, descale_o; + std::shared_ptr descale_dP, descale_dO, descale_dO_t; + std::shared_ptr scale_s, scale_dP; + std::shared_ptr scale_dQ, scale_dK, scale_dV; + std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; + std::shared_ptr seq_q, seq_kv; + std::shared_ptr dropout_seed, dropout_offset; + + // Q, K, V, O, dO, stats, attn_scale + std::vector q_strides(4), k_strides(4), v_strides(4), o_strides(4), dO_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), + k_strides.data(), v_strides.data(), qkv_layout); + generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); + generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_strides.data(), do_format); + Q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q") + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_strides) + .set_data_type(qkv_tensor_type)); + K = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_strides) + .set_data_type(qkv_tensor_type)); + V = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_strides) + .set_data_type(qkv_tensor_type)); + O = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("O") + .set_dim({b, h, s_q, d_v}) + .set_stride(o_strides) + .set_data_type(o_tensor_type)); + dO = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(do_tensor_type)); + Stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Stats") + .set_dim({b, h, s_q, 1}) + .set_stride({h * s_q, s_q, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); + + // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV + if (is_delayed_scaling || is_current_scaling) { + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); + descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); + descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); + scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); + descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); + scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); + if (is_current_scaling && is_O_in_F16) { + descale_o = mha_graph->tensor(1.0f); + } else { + descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); + } + descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); + if (is_delayed_scaling) { + scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); + scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); + scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); + } + if (is_current_scaling) { + scale_dQ = mha_graph->tensor(1.0f); + scale_dK = mha_graph->tensor(1.0f); + scale_dV = mha_graph->tensor(1.0f); + } + } else if (is_mxfp8) { + 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 = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; + NVTE_QKV_Format do_scale_format_ = + (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; + // Q_t, K_t, dO_t, dO_f16 + std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_t_strides.data(), q_format); + generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_t_strides.data(), kv_format); + generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_t_strides.data(), do_format); + Q_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q_t") .set_dim({b, h, s_q, d_qk}) - .set_stride(q_strides) + .set_stride(q_t_strides) .set_data_type(qkv_tensor_type)); - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") + K_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K_t") .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) - .set_data_type(qkv_tensor_type)); - V = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_strides) + .set_stride(k_t_strides) .set_data_type(qkv_tensor_type)); - O = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_strides) - .set_data_type(o_tensor_type)); - dO = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") + dO_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_t") .set_dim({b, h, s_q, d_v}) - .set_stride(dO_strides) + .set_stride(dO_t_strides) .set_data_type(do_tensor_type)); - Stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Stats") - .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_f16") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(o_tensor_type)); + // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), + k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, + q_t_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, + k_t_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, + v_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, + dO_scale_strides.data(), do_scale_format_); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, + dO_t_scale_strides.data(), do_scale_format_); + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_q_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) + .set_stride(q_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k_t") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) + .set_stride(k_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO") + .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) + .set_stride(dO_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) + .set_stride(dO_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + } + + 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 = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } + } + if (is_causal_bottom_right) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } + + // sdpa_backward_options.set_alibi_mask(is_alibi); + + // if (is_bias) { + // bias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("bias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("dBias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // sdpa_backward_options.set_bias(bias); + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + // sdpa_backward_options.set_dbias(dBias); + // } + // } + + if (cudnn_runtime_version >= 91900) { + sdpa_backward_options.set_deterministic_algorithm(deterministic); + } + + if (is_padding) { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV - if (is_delayed_scaling || is_current_scaling) { - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); - descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); - scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); - if (is_current_scaling && is_O_in_F16) { - descale_o = mha_graph->tensor(1.0f); - } else { - descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); - } - descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); - if (is_delayed_scaling) { - scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); - scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); - scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); - } - if (is_current_scaling) { - scale_dQ = mha_graph->tensor(1.0f); - scale_dK = mha_graph->tensor(1.0f); - scale_dV = mha_graph->tensor(1.0f); - } - } else if (is_mxfp8) { - 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 = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; - NVTE_QKV_Format do_scale_format_ = - (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; - // Q_t, K_t, dO_t, dO_f16 - std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_t_strides.data(), q_format); - generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_t_strides.data(), kv_format); - generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_t_strides.data(), do_format); - Q_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q_t") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_t_strides) - .set_data_type(qkv_tensor_type)); - K_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K_t") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_t_strides) - .set_data_type(qkv_tensor_type)); - dO_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO_t") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_t_strides) - .set_data_type(do_tensor_type)); - dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO_f16") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_strides) - .set_data_type(o_tensor_type)); - // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t - auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); - std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), - k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, - q_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, - q_t_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, - k_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, - k_t_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, - v_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, - dO_scale_strides.data(), do_scale_format_); - generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, - dO_t_scale_strides.data(), do_scale_format_); - descale_q = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) - .set_stride(q_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_q_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q_t") - .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) - .set_stride(q_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k") - .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) - .set_stride(k_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k_t") - .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) - .set_stride(k_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_v") - .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) - .set_stride(v_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_dO = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_dO") - .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) - .set_stride(dO_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_dO_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_dO_t") - .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) - .set_stride(dO_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - } + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - 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 = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 92100) { - if (window_size_left != -1) { - sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (window_size_right != -1) { - sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } - } - if (is_causal_bottom_right) { - sdpa_backward_options.set_diagonal_band_right_bound(0); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_sink_token(softmax_offset); + d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("d_softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_dsink_token(d_softmax_offset); + } - // sdpa_backward_options.set_alibi_mask(is_alibi); - - // if (is_bias) { - // bias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("bias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("dBias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // sdpa_backward_options.set_bias(bias); - // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation - // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 - // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { - // sdpa_backward_options.set_dbias(dBias); - // } - // } - - if (cudnn_runtime_version >= 91900) { - sdpa_backward_options.set_deterministic_algorithm(deterministic); - } + std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; + if (is_delayed_scaling || is_current_scaling) { + std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = + std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, + mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, + descale_v, descale_o, descale_dO, descale_s, + descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, + scale_dP, sdpa_backward_options)); + } else if (is_mxfp8) { + std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV) = std::apply( + [](const auto&... elems) { return std::make_tuple(elems...); }, + mha_graph->sdpa_fp8_backward(Q, Q_t, K, K_t, V, O, dO_f16, dO, dO_t, Stats, descale_q, + descale_q_t, descale_k, descale_k_t, descale_v, descale_dO, + descale_dO_t, sdpa_backward_options)); + } + std::vector dq_strides(4), dk_strides(4), dv_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, dq_strides.data(), + dk_strides.data(), dv_strides.data(), dqkv_layout); + dQ->set_output(true) + .set_dim({b, h, s_q, d_qk}) + .set_stride(dq_strides) + .set_data_type(dqkv_tensor_type); + dK->set_output(true) + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(dk_strides) + .set_data_type(dqkv_tensor_type); + dV->set_output(true) + .set_dim({b, hg, s_kv, d_v}) + .set_stride(dv_strides) + .set_data_type(dqkv_tensor_type); + amax_dQ->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + amax_dK->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + amax_dV->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + if (is_delayed_scaling || is_current_scaling) { + amax_dP->set_output(true) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + } - if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_backward_options.set_padding_mask(is_padding) - .set_seq_len_q(seq_q) - .set_seq_len_kv(seq_kv); - } + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // O + std::shared_ptr, // Stats + std::shared_ptr, // dO + std::shared_ptr, // attn_scale + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_o + std::shared_ptr, // descale_dO + std::shared_ptr, // descale_s + std::shared_ptr, // descale_dP + std::shared_ptr, // scale_dQ + std::shared_ptr, // scale_dK + std::shared_ptr, // scale_dV + std::shared_ptr, // scale_s + std::shared_ptr, // scale_dP + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr, // dV + std::shared_ptr, // amax_dQ + std::shared_ptr, // amax_dK + std::shared_ptr, // amax_dV + std::shared_ptr> // amax_dP + key_tensors_tuple = + std::make_tuple(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, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP); + auto mxfp8_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t) + : std::make_tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); + auto softmax_offset_tuple = is_softmax_offset ? std::make_tuple(softmax_offset, d_softmax_offset) + : std::make_tuple(nullptr, nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, + bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); +} - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } +// The backward counterpart of fp8_fwd_cached_graph; see there. +static std::shared_ptr> fp8_bwd_cached_graph( + const FusedAttnConfig& cfg, const Fp8BwdGraphInputs& in, cudnnHandle_t handle) { + static GraphCache cache; + return get_or_build_cached_graph(cache, cfg.make_cache_key(), "bwd", handle, + [&] { return build_sdpa_fp8_bwd_graph(cfg, in); }); +} - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_sink_token(softmax_offset); - d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("d_softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_dsink_token(d_softmax_offset); - } +void fused_attn_fp8_bwd_impl( + 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; - std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; - if (is_delayed_scaling || is_current_scaling) { - std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = - std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, - mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, - descale_v, descale_o, descale_dO, descale_s, - descale_dP, scale_s, scale_dQ, scale_dK, - scale_dV, scale_dP, sdpa_backward_options)); - } else if (is_mxfp8) { - std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV) = std::apply( - [](const auto&... elems) { return std::make_tuple(elems...); }, - mha_graph->sdpa_fp8_backward(Q, Q_t, K, K_t, V, O, dO_f16, dO, dO_t, Stats, descale_q, - descale_q_t, descale_k, descale_k_t, descale_v, descale_dO, - descale_dO_t, sdpa_backward_options)); - } - std::vector dq_strides(4), dk_strides(4), dv_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, dq_strides.data(), - dk_strides.data(), dv_strides.data(), dqkv_layout); - dQ->set_output(true) - .set_dim({b, h, s_q, d_qk}) - .set_stride(dq_strides) - .set_data_type(dqkv_tensor_type); - dK->set_output(true) - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(dk_strides) - .set_data_type(dqkv_tensor_type); - dV->set_output(true) - .set_dim({b, hg, s_kv, d_v}) - .set_stride(dv_strides) - .set_data_type(dqkv_tensor_type); - amax_dQ->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dK->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dV->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - if (is_delayed_scaling || is_current_scaling) { - amax_dP->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - } + // Derived once and handed to the cache, which passes them to the graph build, so that the + // graph and the pointers bound to it below cannot be decided differently. Also where an + // unserviceable configuration is rejected. + const Fp8BwdGraphInputs in = derive_fp8_bwd_graph_inputs(cfg); + const bool is_delayed_scaling = in.is_delayed_scaling; + const bool is_current_scaling = in.is_current_scaling; + const bool is_mxfp8 = in.is_mxfp8; + const bool is_O_in_F16 = in.is_O_in_F16; - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // O - std::shared_ptr, // Stats - std::shared_ptr, // dO - std::shared_ptr, // attn_scale - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_o - std::shared_ptr, // descale_dO - std::shared_ptr, // descale_s - std::shared_ptr, // descale_dP - std::shared_ptr, // scale_dQ - std::shared_ptr, // scale_dK - std::shared_ptr, // scale_dV - std::shared_ptr, // scale_s - std::shared_ptr, // scale_dP - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // amax_dQ - std::shared_ptr, // amax_dK - std::shared_ptr, // amax_dV - std::shared_ptr> // amax_dP - key_tensors_tuple = std::make_tuple( - 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, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP); - auto mxfp8_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t) - : std::make_tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); - auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); - auto softmax_offset_tuple = is_softmax_offset - ? std::make_tuple(softmax_offset, d_softmax_offset) - : std::make_tuple(nullptr, nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : 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})); - }); - 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. - { - std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); - auto inserted = cache.insert({descriptor, return_tuple}); - return inserted.first->second; - } - }; + const int64_t b = static_cast(cfg.batch_size); + const int64_t h = static_cast(cfg.num_attn_heads); + // Not const: bound into the variant pack by address as a pass-by-value graph input. + float scaling_factor = cfg.attn_scale; + const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_padding = cfg.is_padding; + const bool is_dropout = (cfg.dropout != 0.0f); + const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + + try { + auto cache_entry = fp8_bwd_cached_graph(cfg, in, handle); 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, 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, cache_cfg); + dropout_seed, dropout_offset] = cache_entry->tensors; + + // This graph is going to be used, so finish the build the cache deferred. + ensure_plans_built("bwd", *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1147,7 +1187,7 @@ void fused_attn_fp8_bwd_impl( } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } -} // NOLINT(readability/fn_size) +} } // namespace fused_attn @@ -1236,14 +1276,11 @@ 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( - graph_cfg, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, + 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); @@ -1351,18 +1388,15 @@ 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( - 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, + 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); @@ -1383,60 +1417,43 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const } } +// Whether cuDNN can run the FP8 forward graph this config asks for: the empty string if it can, +// otherwise cuDNN's own account of why not, which the backend selector reports to the caller. +// +// The question is answered by deriving the graph's inputs and building the graph, which is +// where every rejection comes from -- there is no separate list of rules to keep in step with +// the builder. The graph goes into the same cache the execution path reads, so the work is not +// thrown away and what was checked is what will run. It stops short of build_plans(), the +// expensive step, which the first execution of the graph does instead; see CachedGraph. +// +// The copy below is made for the sake of one flag; see is_supported_f16_fwd for why that +// assignment is direction selection rather than a restatement of the caller's request. std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; - graph_cfg.is_forward = true; - graph_cfg.derive(); + graph_cfg.check_for_forward_support = true; - size_t workspace_size = 0; try { - fused_attn::fused_attn_fp8_fwd_impl( - graph_cfg, - /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, - /*devPtrSoftmaxOffset=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, - /*devPtrDescaleQ=*/nullptr, /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, - /*devPtrDescaleS=*/nullptr, /*devPtrScaleS=*/nullptr, /*devPtrScaleO=*/nullptr, - /*devPtrAmaxO=*/nullptr, /*devPtrAmaxS=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, - /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, - /*devPtrDropoutOffset=*/nullptr, - /*workspace=*/nullptr, &workspace_size, - /*stream=*/static_cast(0), handle); + const fused_attn::Fp8FwdGraphInputs in = fused_attn::derive_fp8_fwd_graph_inputs(graph_cfg); + fused_attn::fp8_fwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception& e) { - return e.what(); + return fused_attn::refusal_reason(e, "is_supported_fp8_fwd: rejected without a reason."); } catch (...) { return "is_supported_fp8_fwd: unknown failure."; } } +// The backward counterpart of is_supported_fp8_fwd; see there. std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { FusedAttnConfig graph_cfg = cfg; - graph_cfg.is_forward = false; - graph_cfg.derive(); + graph_cfg.check_for_forward_support = false; - size_t workspace_size = 0; try { - fused_attn::fused_attn_fp8_bwd_impl( - graph_cfg, - /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, - /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, - /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, - /*devPtrdSoftmaxOffset=*/nullptr, /*devPtrDescaleQ=*/nullptr, - /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, /*devPtrDescaleO=*/nullptr, - /*devPtrDescaledO=*/nullptr, /*devPtrDescaleS=*/nullptr, /*devPtrDescaledP=*/nullptr, - /*devPtrScaleS=*/nullptr, /*devPtrScaledP=*/nullptr, /*devPtrScaledQ=*/nullptr, - /*devPtrScaledK=*/nullptr, /*devPtrScaledV=*/nullptr, /*devPtrAmaxdP=*/nullptr, - /*devPtrAmaxdQ=*/nullptr, /*devPtrAmaxdK=*/nullptr, /*devPtrAmaxdV=*/nullptr, - /*devPtrQ_t=*/nullptr, /*devPtrK_t=*/nullptr, /*devPtrdO_f16=*/nullptr, - /*devPtrdO_t=*/nullptr, /*devPtrDescaleQ_t=*/nullptr, /*devPtrDescaleK_t=*/nullptr, - /*devPtrDescaledO_t=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, - /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, - /*devPtrDropoutOffset=*/nullptr, - /*workspace=*/nullptr, &workspace_size, - /*stream=*/static_cast(0), handle); + const fused_attn::Fp8BwdGraphInputs in = fused_attn::derive_fp8_bwd_graph_inputs(graph_cfg); + fused_attn::fp8_bwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception& e) { - return e.what(); + return fused_attn::refusal_reason(e, "is_supported_fp8_bwd: rejected without a reason."); } catch (...) { return "is_supported_fp8_bwd: unknown failure."; } diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h new file mode 100644 index 0000000000..83a34d2a95 --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -0,0 +1,242 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// ============================================================================ +// The fused-attention graph cache: what a cache entry is, how one is looked up +// or built, and the frontend calls that make a constructed graph usable. +// +// Each of the four build sites (f16 and fp8, forward and backward) differs only +// in how it constructs its graph and which tensors it hands back. Everything +// after that -- the lookup, the locking, the once-per-entry plan build, the +// support check, and the remembering of what cuDNN refused -- is the same at all +// four, and lives here so it has one definition rather than four copies to keep +// in step. +// +// This header is deliberately not part of utils.h: it needs the cuDNN frontend, +// and utils.h is included by translation units (utils.cu) that otherwise do not. +// ============================================================================ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../common.h" +#include "../cudnn_utils.h" +#include "config_and_params.h" +#include "graph_cache_debug.h" + +namespace transformer_engine { +namespace fused_attn { + +// cuDNN's refusal to run a graph, as opposed to a failure to try. The distinction is what makes +// the negative cache in get_or_build_cached_graph() safe: a refusal is a verdict on the +// configuration and reproducible for a given key, so it can be remembered and replayed, whereas +// a failure that came from the machine's state at that moment (an allocation that did not fit, a +// CUDA error left behind by unrelated work) could well succeed on the next attempt and must not +// be turned into a permanent answer. Only the four adjudicating frontend calls in +// validate_and_check_support() raise this; every other failure keeps its ordinary type and is +// re-attempted the next time the key comes around. +struct UnsupportedGraph : public std::runtime_error { + explicit UnsupportedGraph(const std::string &reason) : std::runtime_error(reason) {} +}; + +// The reason string an is_supported_* helper reports for `e`: its message, or `fallback` if it +// has none. Those helpers signal support by returning the empty string, so a refusal that +// arrives without an explanation would be read as an endorsement and the caller would go on to +// run a graph cuDNN has just declined. Nothing raised through NVTE_ERROR can be empty, since it +// prefixes file and line, but that is a property of our macros rather than of every exception +// that can reach a catch clause, and it is not what the contract should rest on. +inline std::string refusal_reason(const std::exception &e, const char *fallback) { + const char *what = e.what(); + return (what != nullptr && what[0] != '\0') ? std::string(what) : std::string(fallback); +} + +// A graph in the cache, plus the tensor attributes needed to bind runtime pointers to it. +// +// Entries are built only as far as check_support(), which is all it takes to decide whether +// a configuration is supported. build_plans() is the kernel-compilation step and the most +// expensive of the five frontend calls, so a support query stops short of it: the query never +// executes the graph, and many of the keys it builds are never executed by anything. The +// execution path finishes the build instead, the first time the graph is needed to run. +// +// plans_built guards that completion. It has to happen exactly once per entry, because the +// cached graph is shared across threads and build_plans() mutates it in place -- two threads +// reaching the same unfinished entry must not both build it. Keeping the flag inside the entry +// keeps it from drifting away from the graph it describes, and leaves unrelated keys free to +// build concurrently. A build that throws leaves the flag unset, so a later call retries +// rather than executing a graph with no plans. +template +struct CachedGraph { + explicit CachedGraph(GraphAndTensors tensors) : tensors(std::move(tensors)) {} + + GraphAndTensors tensors; + std::once_flag plans_built; +}; + +// One build site's cache. Process-wide rather than per-thread so that a graph is reused +// across threads instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a +// shared plan, and cudnn-frontend >= 1.25.0 has a thread-safe execute(). +// +// Refusals are cached alongside the graphs, under the same keys and the same lock. A support +// query for an unsupported configuration is otherwise the most expensive thing this cache sees: +// it builds the whole graph, spends the four frontend calls, and throws the result away, and it +// does so again on every query, because a rejection left nothing behind to find. `unsupported` +// is what it leaves behind -- cuDNN's own account of the refusal, which is the entire useful +// output of a failed query, so nothing is lost by answering from it. Reasons are short strings +// and there is one per refused key, so this grows far slower than the graphs beside it. +template +struct GraphCache { + std::map>> supported; + std::map unsupported; + std::mutex mutex; // guards both maps +}; + +// Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: +// validate, build_operation_graph, create_execution_plans, check_support. The sequence is +// identical for both passes and both backends, so it is defined once here; `pass` only selects +// which set of stage timers the calls are attributed to. +// +// Support is reported by throwing rather than by a return value. NVTE_CHECK_CUDNN_FE raises +// an exception carrying cuDNN's own explanation of the rejection, and that text is what the +// is_supported_* helpers return as the reason a backend was refused -- so a bool here would +// discard the one thing a support probe exists to produce. Callers that are about to execute +// the graph want the throw as well, since there is nothing useful to do with an unsupported +// graph but fail. +// +// The throw is re-raised as UnsupportedGraph, which is what marks it cacheable. These four calls +// are cuDNN adjudicating a graph it has been handed, so a failure among them is a statement about +// the graph rather than about the moment -- which is the property the negative cache needs, and +// the reason the boundary is drawn here rather than around a wider region. build_plans() and +// execute() sit outside it: they commit real resources and can fail for reasons that have nothing +// to do with the configuration. +// +// build_plans() is left out for a second reason as well: it belongs to whoever executes the graph, +// once, the first time it is needed. See CachedGraph. +inline void validate_and_check_support(const char *pass, cudnn_frontend::graph::Graph &graph, + cudnnHandle_t handle) { + try { + graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(graph.validate()); }); + graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::BuildOpGraph, + [&] { NVTE_CHECK_CUDNN_FE(graph.build_operation_graph(handle)); }); + graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(graph.create_execution_plans({cudnn_frontend::HeurMode_t::A})); + }); + graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(graph.check_support()); }); + } catch (const std::exception &e) { + throw UnsupportedGraph(e.what()); + } +} + +// The cached entry for `key`, building and inserting it via `build` if absent. Throws +// UnsupportedGraph if cuDNN refuses the graph -- this time or on an earlier call, the two being +// indistinguishable to the caller by design. +// +// `build` only constructs a graph; this is what puts it through validate_and_check_support(), so +// the entries in the cache are exactly the graphs cuDNN has agreed to run. Those four calls sit +// on the miss path because they are part of building an entry rather than reading one: repeating +// them on a hit would redo the operation graph and the plan search for a graph that has already +// been through both. +// +// `key` must be a normalized key -- FusedAttnConfig::make_cache_key()'s output -- and not a +// raw execution config. Two configs that differ only in a field no graph reads (attn_scale, +// say) have to reach the same entry, which is what normalization is for; passing the raw +// config instead silently multiplies the cache by fields the graph never consumes. +// +// Only the map operations are locked, not `build`. A graph build is the expensive part and +// holding the lock across it would serialize builds of unrelated keys, so two threads racing +// on the same key may both build. That is a wasted build, not a correctness problem: the +// loser drops its own graph and takes the winner's, so every caller of a given key gets one +// shared entry and the once-flag inside it still governs the plan build. The wasted build is +// visible in diagnostics as a BUILD with no matching MISS of its own. The same race on a +// refused key is equally harmless, both threads storing the same reason. +template +std::shared_ptr> get_or_build_cached_graph( + GraphCache &cache, const FusedAttnConfig &key, const char *pass, + cudnnHandle_t handle, BuildFn &&build) { + using Entry = CachedGraph; + + std::shared_ptr cached; + bool refused = false; + std::string reason; + { + std::lock_guard lock(cache.mutex); + auto it = cache.supported.find(key); + if (it != cache.supported.end()) { + cached = it->second; + } else { + auto refusal = cache.unsupported.find(key); + refused = (refusal != cache.unsupported.end()); + if (refused) reason = refusal->second; + } + } + using graph_cache_debug::LookupResult; + LookupResult outcome = LookupResult::Miss; + if (cached != nullptr) { + outcome = LookupResult::Hit; + } else if (refused) { + outcome = LookupResult::Unsupported; + } + graph_cache_debug::record_cache_lookup(pass, outcome, key); + + if (cached != nullptr) return cached; + // Raised rather than returned so that a replayed refusal is the same event as a fresh one: + // every caller already has to handle the build refusing, and none of them would have anything + // else to do with a second, quieter way of saying so. + if (refused) throw UnsupportedGraph(reason); + + std::shared_ptr entry; + try { + entry = std::make_shared(build()); + // Every site's tensor tuple leads with its graph, which is the one thing all four have in + // common and the only element this needs. A tuple that stopped leading with it would fail to + // compile here rather than quietly validate the wrong object. + validate_and_check_support(pass, *std::get<0>(entry->tensors), handle); + } catch (const UnsupportedGraph &e) { + { + std::lock_guard lock(cache.mutex); + cache.unsupported.insert({key, e.what()}); + } + graph_cache_debug::record_unsupported(pass); + throw; + } + graph_cache_debug::record_build(pass); + { + std::lock_guard lock(cache.mutex); + return cache.supported.insert({key, std::move(entry)}).first->second; + } +} + +// Runs the plan build that get_or_build_cached_graph() left undone, once per entry. +// +// Call this only when the graph is about to be executed, which is why it is a separate step +// rather than the tail of the lookup: a support query builds entries that nothing ever runs, and +// kernel compilation is the most expensive of the five frontend calls, so a query that paid for +// it would be paying for nothing. See CachedGraph for why the flag lives inside the entry and +// what a throw here leaves behind. +template +void ensure_plans_built(const char *pass, CachedGraph &entry) { + std::call_once(entry.plans_built, [&] { + cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.tensors); + graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); + graph_cache_debug::record_plans_built(pass); + }); +} + +} // namespace fused_attn +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 033edf2ead..a23812dc67 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -7,14 +7,34 @@ // ============================================================================ // 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. +// Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG. Two verbosity levels: +// =1 : low volume. Cache event counters, a BUILD and a PLANS line per build, an +// UNSUP line per configuration cuDNN refuses, and the end-of-run SUMMARY +// (aggregate + per-thread) and stage timings. This is enough to diagnose +// redundant rebuilds and profile build cost. +// =2 : high volume (trace). Additionally emits a per-lookup HIT/MISS/NOSUP line +// with the full shorthand cache key and a per-execution EXEC line. Use only +// when you need to see *which* shapes are hitting/missing -- these fire on +// every cache lookup and execution, so at suite scale they add I/O and +// serialize threads on the stderr lock. No timed region writes to stderr, so +// the stage timings stay sound, but they are measured under more contention +// than at level 1 and read a little high. +// +// NOSUP is a hit on the negative cache: a key cuDNN has already refused, answered +// from the stored refusal instead of by building the graph again. +// +// An optional ":" suffix picks which processes emit, defaulting to rank 0 +// so that output does not scale with the world size: "1:all" for every rank, +// "2:0,3" for a specific set. See `rank_selected` for when overriding pays off. // ============================================================================ #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 @@ -23,6 +43,10 @@ #include #include #include +#include +#include +#include +#include #include "config_and_params.h" @@ -30,30 +54,106 @@ 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 = [] { +// Rank of this process as reported by the launcher, or -1 when there is no +// launcher (a single-process run). First variable that is set wins. +inline int launcher_rank() { + static const int rank = []() -> int { + for (const char *var : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { + const char *v = std::getenv(var); + if (v != nullptr && v[0] != '\0') return std::atoi(v); + } + return -1; + }(); + return rank; +} + +// Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=default, +// 2=trace). Single read at startup, cached. Negligible overhead when unset. +inline int debug_level() { + static const int lvl = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; + const int v = std::atoi(e); // stops at the optional ":" suffix + return v > 0 ? v : 1; // any non-empty, non-"0" value enables at least level 1 + }(); + return lvl; +} + +// Whether this process emits diagnostics. Every rank writes to the same stderr, +// so emitting from all of them multiplies the volume by the world size -- and +// under data/tensor parallelism the ranks are running identical shapes, so the +// copies say the same thing. Hence rank 0 only by default. +// +// Context parallelism is the case worth overriding for: the ranks run different +// subsets of the per-step regimes (under p2p, rank 0 never sees the lower-triangle +// config that the last rank does), so their build counts genuinely differ. +// Select with the ":" suffix, e.g. "1:all" or "2:0,3". +inline bool rank_selected() { + static const bool selected = [] { + const int rank = launcher_rank(); + if (rank < 0) return true; // sole process, nothing to filter const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - return e != nullptr && e[0] != '\0' && e[0] != '0'; + const char *sep = (e != nullptr) ? std::strchr(e, ':') : nullptr; + if (sep == nullptr) return rank == 0; + const std::string list(sep + 1); + if (list == "all") return true; + for (size_t pos = 0; pos <= list.size();) { + const size_t comma = list.find(',', pos); + const std::string tok = + list.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + if (!tok.empty() && std::atoi(tok.c_str()) == rank) return true; + if (comma == std::string::npos) break; + pos = comma + 1; + } + return false; }(); - return on; + return selected; } -// More readable, shorter thread IDs (0, 1, 2, ...). +// Diagnostics are on at level >= 1, and only for the selected ranks. Unselected +// ranks skip the counters too, so they pay nothing beyond this check. +inline bool enabled() { return debug_level() >= 1 && rank_selected(); } + +// Per-lookup / per-exec trace lines are gated behind level >= 2. +inline bool trace_enabled() { return debug_level() >= 2; } + +// Identifies the emitting process. Distributed PyTorch runs one process per rank +// and they all share this stderr, so without this every line would be ambiguous +// (thread ids restart at 0 in each process). Rank comes from the launcher, if any. +inline const std::string &process_tag() { + static const std::string *tag = [] { + auto *s = new std::string("pid=" + std::to_string(static_cast(::getpid()))); + if (launcher_rank() >= 0) *s += " rank=" + std::to_string(launcher_rank()); + return s; + }(); + return *tag; +} + +// More readable, shorter thread IDs (0, 1, 2, ...). These are assignment order, +// not identity: tid=0 is whichever thread touched this cache first. The one-shot +// THREAD line below maps them to OS thread ids for correlating with nsys/gdb. 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; } +// OS-level thread id, as reported by nsys/gdb/`top -H`. Via syscall rather than +// gettid() so this does not require glibc >= 2.30. +inline int64_t os_thread_id() { return static_cast(::syscall(SYS_gettid)); } + // 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 +// - BUILD: a graph built and cached in response to a cache miss. Built only as far as +// check_support(), which is all a support probe needs. +// - PLANS: a cached graph finished with build_plans(), the kernel compilation that the +// BUILD above deferred. At most one per BUILD, on the first execution of that +// graph, so BUILD minus PLANS is how many graphs were built for a support probe +// and never used to run anything. // - 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 @@ -63,9 +163,11 @@ inline void register_summary_once(); struct EventCounters { std::atomic built{0}; + std::atomic plans{0}; std::atomic exec{0}; std::atomic hit{0}; std::atomic miss{0}; + std::atomic unsup{0}; }; inline EventCounters &counters(bool is_fwd) { @@ -74,45 +176,182 @@ inline EventCounters &counters(bool is_fwd) { 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=%" 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)); +// Per-thread counters, so the summary can break down build/exec/hit/miss by +// thread. In the single-process context-parallel case each device is driven by +// its own thread, so this reveals which thread built/executed what. +struct ThreadCounters { + unsigned tid = 0; + EventCounters fwd; + EventCounters bwd; +}; + +// The registry and its mutex are heap-allocated and deliberately never freed. +// Function-local static destructors and atexit handlers run as a single sequence, +// in reverse order of construction/registration. This registry is built lazily, so +// it can be constructed *after* the summary handler is registered -- in which case +// it would be destroyed *before* that handler runs, leaving the handler to lock a +// destroyed mutex and walk a destroyed vector. Leaking removes the ordering +// question rather than reasoning about it, and the cost is bounded: one mutex and +// one vector for the process, reclaimed by the OS at exit anyway. +inline std::mutex &thread_registry_mutex() { + static std::mutex *m = new std::mutex(); + return *m; +} +inline std::vector &thread_registry() { + static std::vector *v = new std::vector(); + return *v; +} + +// This thread's counter block, leaked for a related but distinct reason: a worker +// thread can exit long before the process does, while the registry keeps a pointer +// to its block for the end-of-run summary. Tying the block's lifetime to the +// thread would leave that pointer dangling. One small struct per thread. +inline ThreadCounters &thread_counters() { + static thread_local ThreadCounters *tc = [] { + auto *p = new ThreadCounters(); + p->tid = thread_seq_id(); + { + std::lock_guard lock(thread_registry_mutex()); + thread_registry().push_back(p); + } + // One line per thread, mapping the short id to something nsys/gdb can match. + std::fprintf(stderr, "[FUSED-ATTN-CACHE] %s | THREAD | tid=%-3u os_tid=%" PRId64 "\n", + process_tag().c_str(), p->tid, os_thread_id()); + std::fflush(stderr); + return p; + }(); + return *tc; +} + +inline EventCounters &thread_counters(bool is_fwd) { + ThreadCounters &tc = thread_counters(); + return is_fwd ? tc.fwd : tc.bwd; +} + +// Format one counter block (aggregate or a single thread's) as one line. +// `tid_field` is the whole thread column, e.g. "tid=3"; the aggregate row passes +// "tid=all" so that it cannot be misread as thread 0's row. +// +// The columns are meant to be read against two identities. Every lookup lands in exactly one of +// miss and hit, and every miss ends in exactly one of built and unsup -- so `miss = built + unsup` +// and a shortfall in either means a build died of something other than a refusal. `built >= plans` +// always, the difference being graphs that a support query built and nothing has yet run. +inline std::string format_counter_line(const char *event, const char *tid_field, + const EventCounters &f, const EventCounters &b) { + char buf[768]; + std::snprintf(buf, sizeof(buf), + "[FUSED-ATTN-CACHE] %s | %-11s | %-7s | fwd miss=%4" PRIu64 ", hit=%4" PRIu64 + ", built=%4" PRIu64 ", unsup=%4" PRIu64 ", plans=%4" PRIu64 ", exec=%4" PRIu64 + " | bwd miss=%4" PRIu64 ", hit=%4" PRIu64 ", built=%4" PRIu64 ", unsup=%4" PRIu64 + ", plans=%4" PRIu64 ", exec=%4" PRIu64 "\n", + process_tag().c_str(), event, tid_field, f.miss.load(std::memory_order_relaxed), + f.hit.load(std::memory_order_relaxed), f.built.load(std::memory_order_relaxed), + f.unsup.load(std::memory_order_relaxed), f.plans.load(std::memory_order_relaxed), + f.exec.load(std::memory_order_relaxed), b.miss.load(std::memory_order_relaxed), + b.hit.load(std::memory_order_relaxed), b.built.load(std::memory_order_relaxed), + b.unsup.load(std::memory_order_relaxed), b.plans.load(std::memory_order_relaxed), + b.exec.load(std::memory_order_relaxed)); + return std::string(buf); +} + +inline void print_counter_block(const char *event, const char *tid_field, const EventCounters &f, + const EventCounters &b) { + const std::string line = format_counter_line(event, tid_field, f, b); + std::fputs(line.c_str(), stderr); std::fflush(stderr); } +inline void print_counters(const char *event) { + char tid_field[16]; + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); + print_counter_block(event, tid_field, counters(/*is_fwd=*/true), counters(/*is_fwd=*/false)); +} + +// A graph built through check_support() and cached. Call after the build, from the miss +// path that performed it. 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); + thread_counters(is_fwd).built.fetch_add(1, std::memory_order_relaxed); print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD"); } +// The build_plans() a BUILD deferred, now completed. Call from inside the std::call_once +// that runs it, after the call returns rather than before: build_plans() throws without +// setting the once_flag, leaving a later execution to retry it, so counting on the way out +// keeps this a count of graphs that reached a runnable state. Like BUILD this fires once +// per distinct cache key, so it stays on the level-1 path. +inline void record_plans_built(const char *pass) { + if (!enabled()) return; + register_summary_once(); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + counters(is_fwd).plans.fetch_add(1, std::memory_order_relaxed); + thread_counters(is_fwd).plans.fetch_add(1, std::memory_order_relaxed); + print_counters(is_fwd ? "fwd PLANS" : "bwd PLANS"); +} + +// A build that cuDNN refused, now remembered as a negative cache entry. Call from the miss path +// that attempted it, in place of record_build(): a refusal and a build are the two ways a miss +// can end, and counting both keeps `miss = built + unsup` true. Fires once per distinct refused +// key -- the second query for that key is a hit -- so it stays on the level-1 path. +inline void record_unsupported(const char *pass) { + if (!enabled()) return; + register_summary_once(); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + counters(is_fwd).unsup.fetch_add(1, std::memory_order_relaxed); + thread_counters(is_fwd).unsup.fetch_add(1, std::memory_order_relaxed); + print_counters(is_fwd ? "fwd UNSUP" : "bwd UNSUP"); +} + 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); + thread_counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); + // The per-exec line fires on every execution; keep it out of the level-1 path. + if (!trace_enabled()) return; print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC"); } -inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { +// What a lookup found. Unsupported is the negative-cache case: a key whose graph cuDNN has +// already refused, so the answer is a remembered refusal rather than a graph. +enum class LookupResult { Miss, Hit, Unsupported }; + +// `key` is the normalized cache key -- make_cache_key()'s output, the exact value the +// lookup was performed with -- not the execution config it was derived from. That is +// deliberate: HIT/MISS is decided by comparing keys, so a trace of anything else cannot +// explain its own outcome. Logging the pre-normalization config would show pairs of +// identical lines with opposite outcomes (normalization having collapsed a difference, +// e.g. bottom_right_diagonal or the THD token counts) and pairs of differing lines that +// both hit (the difference being in a field the key drops, e.g. attn_scale). Diffing two +// MISS lines here instead names exactly the fields responsible for the extra build. +// +// The cost is that fields normalization overwrites are no longer visible in their +// original form: attn_scale reads 1, ragged num_tokens read 0, and max_seqlen/batch_size +// read their bucketed values. Recover those from the caller if a line needs to be traced +// back to a specific test case. +inline void record_cache_lookup(const char *pass, LookupResult result, const FusedAttnConfig &key) { if (!enabled()) return; register_summary_once(); - EventCounters &pc = counters(std::strcmp(pass, "fwd") == 0); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + // Unsupported counts as a hit: what the hit column measures is lookups that were answered + // without building anything, and a remembered refusal is one of those. Which kind of answer + // it was shows in the trace line, and the running total of refusals is the unsup column. + const bool hit = (result != LookupResult::Miss); + EventCounters &pc = counters(is_fwd); (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); + EventCounters &tpc = thread_counters(is_fwd); + (hit ? tpc.hit : tpc.miss).fetch_add(1, std::memory_order_relaxed); + // The per-lookup config dump is the highest-volume line (one per cache lookup); + // keep it out of the level-1 path and off the stderr lock unless tracing. + if (!trace_enabled()) return; std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %-3s %-4s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%d " + "[FUSED-ATTN-CACHE] %s | %-3s %-5s | tid=%u dev=%d | 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 @@ -121,46 +360,74 @@ inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfi " 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)); - std::fflush(stderr); + process_tag().c_str(), pass, + result == LookupResult::Miss ? "MISS" : (result == LookupResult::Hit ? "HIT" : "NOSUP"), + thread_seq_id(), key.device_id, static_cast(key.is_training), + static_cast(key.deterministic), static_cast(key.cuda_graph), + static_cast(key.return_max_logit), static_cast(key.check_for_forward_support), + static_cast(key.attn_mask_type), static_cast(key.bias_type), + static_cast(key.window_size_left), static_cast(key.window_size_right), + static_cast(key.bottom_right_diagonal), static_cast(key.softmax_type), + static_cast(key.scaling_mode), static_cast(key.dropout), + static_cast(key.attn_scale), static_cast(key.qkv_dtype), + static_cast(key.o_dtype), static_cast(key.do_dtype), + static_cast(key.dqkv_dtype), static_cast(key.qkv_layout), + static_cast(key.o_format), static_cast(key.do_format), + static_cast(key.dqkv_layout), static_cast(key.qkv_scale_inv_format), + static_cast(key.do_scale_inv_format), static_cast(key.batch_size), + static_cast(key.num_attn_heads), static_cast(key.num_gqa_groups), + static_cast(key.head_dim_qk), static_cast(key.head_dim_v), + static_cast(key.max_seqlen_q), static_cast(key.max_seqlen_kv), + static_cast(key.num_tokens_q), static_cast(key.num_tokens_kv), + static_cast(key.bucketed_batch_size), + static_cast(key.bucketed_num_tokens_q), + static_cast(key.bucketed_num_tokens_kv), static_cast(key.num_pages_k), + static_cast(key.num_pages_v), static_cast(key.page_size_k), + static_cast(key.page_size_v), static_cast(key.max_pages_per_seq_k), + static_cast(key.max_pages_per_seq_v), static_cast(key.bias_batch_size), + static_cast(key.bias_num_heads), static_cast(key.bias_seqlen_q), + static_cast(key.bias_seqlen_kv)); } // ============================================================================ -// Graph build timings for individual cuDNN-frontend calls in forward/backward: -// e.g. `validate`, `build_operation_graph`, `create_execution_plans`, -// `check_support`, `build_plans` +// Graph build timings. +// +// A cuDNN graph build is a fixed sequence of frontend calls, and which one +// dominates determines what to do about a slow build: time in `check_support` +// and `build_plans` is heuristic selection and kernel compilation, largely +// intrinsic to the shape, whereas time in `validate` or `build_operation_graph` +// is graph-construction cost on our side of the boundary. Timing the stages +// separately is what makes that distinction; one duration per build cannot. +// +// Each stage is wrapped where it is called -- graph_cache.h, which is where all +// five frontend calls live -- and accumulates into the table below, under the +// pass its caller was serving. +// The end-of-run summary reports each as a mean over its calls. Only sums are +// kept, so the mean is all that can be recovered -- and since a build happens +// once per distinct cache key, those calls span different shapes rather than +// repeating one. Read a stage mean as where build time goes in aggregate, not as +// the cost of any particular build. // ============================================================================ +// The frontend calls that make up a build, in the order they run. `kCount` must +// stay last: it sizes the table below. `kStageNames` is indexed by these values +// when the summary prints, so the two must be kept in the same order. 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"}; + +// Totals for one (pass, stage) pair. Relaxed ordering is sufficient: these +// counters order nothing, and the only read happens once, after the threads that +// wrote them are done. struct StageTiming { std::atomic calls{0}; std::atomic time_ns{0}; }; + +// Bucketed by pass, so the summary can report the cost of each stage separately +// for forward and backward. Unlike the thread registry above, this table needs no +// leak to outlive the exit handler that reads it: it holds nothing but atomics, so +// it is trivially destructible and no destructor is registered for it at all. constexpr size_t kStageBuckets = 2 * static_cast(BuildStage::kCount); inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { static std::array table{}; @@ -169,6 +436,14 @@ inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { return table[idx]; } +// Times one stage: clock read in the constructor, accumulated in the destructor. +// Recording on scope exit rather than at an explicit stop() keeps a failing stage +// measurable -- the frontend calls are wrapped in NVTE_CHECK_CUDNN_FE, which +// throws, and the destructor still runs during unwinding -- so a build that dies +// in `check_support` contributes its time to failure instead of vanishing from the +// summary. `on` is latched at construction rather than re-tested in the destructor, +// which is what keeps that symmetric: the destructor can never accumulate against a +// `start` the constructor left unset. struct ScopedBuildTimer { BuildStage stage; bool on; @@ -191,6 +466,11 @@ struct ScopedBuildTimer { } }; +// Time `fn` as `stage` of the given pass ("fwd"/"bwd", matching the record_* +// helpers above). Preferred over declaring a ScopedBuildTimer at the call site: +// the measured region is exactly the call passed in, so surrounding work cannot +// drift into it as that code changes. With diagnostics off this costs the pass +// comparison and one cached-flag check; both are per build, not per lookup. template inline void timer(const char *pass, BuildStage stage, Fn &&fn) { ScopedBuildTimer scoped(std::strcmp(pass, "fwd") == 0, stage); @@ -204,7 +484,27 @@ inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; - print_counters("SUMMARY"); + // Build the whole summary in memory and emit it with a single write, so + // that the blocks of concurrently-exiting processes (one per rank under + // torchrun) stay grouped instead of interleaving line by line. + std::string block; + block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary begin =====\n"; + // Per-thread breakdown (sorted by tid). Useful in the single-process + // context-parallel case where each device runs on its own thread. + { + std::lock_guard lock(thread_registry_mutex()); + std::vector blocks = thread_registry(); + std::sort(blocks.begin(), blocks.end(), + [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); + for (const ThreadCounters *tc : blocks) { + char tid_field[16]; + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); + block += format_counter_line("SUMMARY-TID", tid_field, tc->fwd, tc->bwd); + } + } + // Totals last, so they read as the sum of the per-thread lines above. + block += format_counter_line("SUMMARY", "tid=all", counters(/*is_fwd=*/true), + counters(/*is_fwd=*/false)); for (int p = 0; p < 2; ++p) { const bool is_fwd = (p == 0); const char *pass = is_fwd ? "fwd" : "bwd"; @@ -215,12 +515,16 @@ 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=%" PRIu64 - " | time=%9.1f ms | avg=%9.3f ms/call\n", - pass, kStageNames[i], n, total_ms, total_ms / n); + char line[288]; + std::snprintf(line, sizeof(line), + "[FUSED-ATTN-CACHE] %s | %-3s %-22s | calls=%" PRIu64 + " | time=%9.3f ms/call\n", + process_tag().c_str(), pass, kStageNames[i], n, total_ms / n); + block += line; } } + block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary end =====\n"; + std::fwrite(block.data(), 1, block.size(), stderr); std::fflush(stderr); }); return true; diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index e240e2a421..391a22768a 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -8,7 +8,6 @@ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ #include -#include #include "../common.h" #include "transformer_engine/fused_attn.h" From 6d2b5776ae66e3438f76fbafb19371f97b14731b Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 08:09:42 -0700 Subject: [PATCH 76/88] revert consolidating of derive() and make nvte_get_backend_v2 only read and not modify cfg Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.h | 22 ++++--- .../common/fused_attn/fused_attn.cpp | 66 ++++++++++++------- 2 files changed, 55 insertions(+), 33 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 5d04499409..a9c65230f8 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -182,8 +182,14 @@ struct FusedAttnConfig { // Derive fields such as bucketed batch_size or num_tokens for THD, based on input fields // that have been set by the caller. Call once, after the last input field is set and before // the config reaches a graph build, a cache lookup, or a support query -- all of which read - // derived fields. nvte_get_fused_attn_backend_v2() is where that happens for every config - // that enters this library, so nothing downstream of it needs to derive again. + // derived fields. + // + // Called by whoever owns the config, at the point it stops being edited: the execution entry + // points (nvte_fused_attn_fwd_v2 and its backward counterpart) on the config they go on to run, + // and nvte_get_fused_attn_backend_v2() on a copy of the caller's, so that asking whether a + // configuration is supported does not modify it. Nothing further in is expected to derive + // again, and check_derived() is what holds them to that. Idempotent, so a config that is + // derived and then re-derived is unharmed. void derive(); // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. @@ -196,14 +202,14 @@ struct FusedAttnConfig { // Assert that `cfg` has been through derive(), for code about to read a derived field. Worth // asserting rather than assuming because the failure is silent: an unset bucketed_batch_size or // q_format reads as zero, which is a legal value that yields a graph of the wrong shape and a -// key that collides with unrelated configs. Deriving happens in exactly one place -// (nvte_get_fused_attn_backend_v2), so this is what keeps a new path into the builders from +// key that collides with unrelated configs. Deriving happens at the library's entry points rather +// than here, where it would be needed, so this is what keeps a new path into the builders from // quietly skipping it. inline void check_derived(const FusedAttnConfig &cfg) { - NVTE_CHECK( - cfg.is_derived, - "FusedAttnConfig reached a graph build with its derived fields unset. Every config " - "must pass through FusedAttnConfig::derive() first; see nvte_get_fused_attn_backend_v2."); + NVTE_CHECK(cfg.is_derived, + "FusedAttnConfig reached a graph build with its derived fields unset. Every config " + "must pass through FusedAttnConfig::derive() first; see the entry points in " + "fused_attn.cpp."); } 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 5ef661ffe0..f7dacb9cdc 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 { -// The per-thread storage for the diagnostic string; it's re-used (cleared + re-populated) +// The per-thread storage for the diagnostic string; it is 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; @@ -240,6 +240,15 @@ void set_message(const char **message, std::string reason) { *message = fused_attn_backend_message_buffer.c_str(); } +// Records `reason` if `rejected`, and reports it, so that an early rejection reads as the one +// statement it is: `if (set_message_if(cond, message, "why")) return NVTE_No_Backend;`. The +// reason is built whether or not it is used, which is why it stays a plain string here -- these +// are short literals on a path that goes on to build cuDNN graphs. +bool set_message_if(bool rejected, const char **message, std::string reason) { + if (rejected) set_message(message, std::move(reason)); + return rejected; +} + } // namespace // select a backend for fused attention; the diagnostic message is based on the first failure, not cumulative. @@ -247,21 +256,22 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi const char **message) { using namespace transformer_engine; using namespace transformer_engine::fused_attn; - // Every config entering this library passes through here on its way to a graph, so this is the - // one place that has to fill the derived fields, and it does so in place. The caller keeps - // ownership; what it gets back is its own config with the blanks filled in, which is what the - // execution path then hands to the backend it selected -- nvte_fused_attn_fwd_v2() queries with - // the very config it goes on to run, so deriving here is what lets the run reuse the graph the - // query built rather than key a second one. Deriving is idempotent, so a config that arrives - // already derived is unharmed. It is a write, though, so one config object must not be queried - // from two threads at once; every caller here builds its config as a local, one per call. - FusedAttnConfig &cfg = *get_fused_attn_config_mutable(config); + // Derived on a copy, leaving the caller's config untouched: this function answers a question + // about a configuration and has no business editing one, and a query that wrote to its argument + // could not be asked about the same config from two threads at once. The copy costs nothing that + // matters here, since deriving is a version check and some arithmetic. + // + // The execution path derives its own config before calling this (see nvte_fused_attn_fwd_v2), + // and still reuses whatever graph the query builds: both derive the same fields from the same + // inputs, so make_cache_key() lands on the same entry. Deriving is idempotent, so re-deriving + // an already-derived config here changes nothing. + FusedAttnConfig cfg = *get_fused_attn_config(config); cfg.derive(); set_message(message, ""); cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(cfg.qkv_layout); - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(cfg.qkv_layout); + const auto qkv_format = nvte_get_qkv_format(cfg.qkv_layout); + const auto 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 @@ -270,26 +280,26 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi fused_attn::get_ragged_offset_dtype(layout_group, cfg.num_attn_heads, cfg.num_gqa_groups, cfg.max_seqlen_q, cfg.max_seqlen_kv, cfg.head_dim_qk, cfg.head_dim_v) == DType::kInt64); - if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { - set_message(message, - "Configuration requires 64-bit ragged offsets, which require " - "cuDNN >= 9.5."); + if (set_message_if(requires_64bit_ragged_offset && cudnn_runtime_version < 90500, message, + "Configuration requires 64-bit ragged offsets, which require " + "cuDNN >= 9.5.")) { return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } // THD requires padding-style mask - if (qkv_format == NVTE_QKV_Format::NVTE_THD && - cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && - cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && - cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { - set_message(message, - "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + if (set_message_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, + message, + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask.")) { 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."); + // TE's cuDNN fused-attention graph does not represent pre-scale bias. + if (set_message_if(cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS, message, + "Fused attention does not support pre-scale bias.")) { return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -427,6 +437,9 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); FusedAttnConfig cfg = p.make_config(); + // Derived here, not by the query below: the query works on its own copy, and it is this config + // that goes on to the backend and must arrive with its derived fields filled in. + cfg.derive(); 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); @@ -529,6 +542,9 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); FusedAttnConfig cfg = p.make_config(); + // Derived here, not by the query below: the query works on its own copy, and it is this config + // that goes on to the backend and must arrive with its derived fields filled in. + cfg.derive(); 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); From ce58a9942e007479bed0cf8e392980572c4edeff Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 08:26:28 -0700 Subject: [PATCH 77/88] incorporate PR5 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/config_and_params.cpp | 2 +- .../common/fused_attn/config_and_params.h | 9 +++ .../fused_attn_f16_arbitrary_seqlen.cu | 56 ++++++++++++------- .../attention/dot_product_attention/utils.py | 6 +- 4 files changed, 49 insertions(+), 24 deletions(-) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 91b4445617..f42850ddad 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -114,7 +114,7 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { 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 (supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_)) { if (cache_cfg.is_ragged_q) { cache_cfg.max_seqlen_q = cache_cfg.bucketed_num_tokens_q; } diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index a9c65230f8..3bbf77dfa3 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -19,6 +19,15 @@ namespace transformer_engine { namespace fused_attn { +// Whether a ragged (THD) graph can be built at packed token-count dimensions with ragged +// Stats/LSE. SM8x and SM120 require dense, BHSD-like dimensions at max_seqlen for the auxiliary +// tensors instead. Graph construction, auxiliary-tensor allocation and make_cache_key() all +// answer this question, and a disagreement between them would key a graph by dimensions it was +// not built with, so they share this one definition. +inline constexpr bool supports_packed_ragged_graph(size_t cudnn_runtime_version, int sm_arch) { + return cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; +} + struct FusedAttnConfig { // basic attention settings bool is_training = true; 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 b07ec848ca..da3bf7a875 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 @@ -64,6 +64,10 @@ struct F16FwdGraphInputs { // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever // the bucketing above did to `b`. int64_t actual_b; + // Whether this architecture takes the packed token-count representation above; see + // supports_packed_ragged_graph. Carried here so the graph build reads the same answer the + // dimensions were derived from rather than querying the device again. + bool use_packed_ragged_graph; bool use_ragged_stats; DType ragged_offset_type; RaggedOffsetMultipliers offset_mults; @@ -86,6 +90,9 @@ static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + const bool use_packed_ragged_graph = + supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); + int64_t b = static_cast(cfg.batch_size); int64_t s_q = static_cast(cfg.max_seqlen_q); int64_t s_kv = static_cast(cfg.max_seqlen_kv); @@ -93,10 +100,10 @@ static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) const int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] - // as interleaved and rejects them. Use BHSD-like dimensions/strides with max_seqlen at plan build - // so the check passes; ragged offset still provides variable-length boundaries. - if (sm_arch_ != 120) { + // SM8x and SM120 need dense, BHSD-like dimensions/strides at max_seqlen: on SM120 the cuDNN + // support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] as interleaved and + // rejects them. The ragged offsets still provide the variable-length boundaries either way. + if (use_packed_ragged_graph) { // 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. // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: @@ -110,7 +117,7 @@ static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) } } - const bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; + const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph; const DType ragged_offset_type = use_cu_seqlens_directly ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. @@ -124,7 +131,14 @@ static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) // Field order must match F16FwdGraphInputs; one per line so that it can be checked by eye. return F16FwdGraphInputs{ - b, s_q, s_kv, actual_b, use_ragged_stats, ragged_offset_type, offset_mults, + b, + s_q, + s_kv, + actual_b, + use_packed_ragged_graph, + use_ragged_stats, + ragged_offset_type, + offset_mults, }; } @@ -675,6 +689,7 @@ struct F16BwdGraphInputs { int64_t s_q; int64_t s_kv; int64_t actual_b; + bool use_packed_ragged_graph; bool use_ragged_stats; DType ragged_offset_type; }; @@ -689,6 +704,9 @@ static F16BwdGraphInputs derive_f16_bwd_graph_inputs(const FusedAttnConfig &cfg) const auto cudnn_runtime_version = cudnnGetVersion(); const int sm_arch_ = cuda::sm_arch(cuda::current_device()); + const bool use_packed_ragged_graph = + supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); + int64_t b = static_cast(cfg.batch_size); int64_t s_q = static_cast(cfg.max_seqlen_q); int64_t s_kv = static_cast(cfg.max_seqlen_kv); @@ -696,8 +714,8 @@ static F16BwdGraphInputs derive_f16_bwd_graph_inputs(const FusedAttnConfig &cfg) const int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). - if (sm_arch_ != 120) { + // SM8x and SM120 require dense, BHSD-like strides at max_seqlen (see fwd). + if (use_packed_ragged_graph) { // 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. // The batch is bucketed unconditionally here, where the forward pass guards it: only @@ -711,14 +729,14 @@ static F16BwdGraphInputs derive_f16_bwd_graph_inputs(const FusedAttnConfig &cfg) } } - const bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; + const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph; // 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; // Field order must match F16BwdGraphInputs; one per line so that it can be checked by eye. return F16BwdGraphInputs{ - b, s_q, s_kv, actual_b, use_ragged_stats, ragged_offset_type, + b, s_q, s_kv, actual_b, use_packed_ragged_graph, use_ragged_stats, ragged_offset_type, }; } @@ -762,7 +780,7 @@ static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); - const int sm_arch_ = cuda::sm_arch(cuda::current_device()); + const bool use_packed_ragged_graph = in.use_packed_ragged_graph; const bool use_ragged_stats = in.use_ragged_stats; const DType ragged_offset_type = in.ragged_offset_type; @@ -865,7 +883,7 @@ static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig if (use_ragged_stats) { sdpa_backward_options.set_max_total_seq_len_q(s_q); } - if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + if (is_ragged_kv && use_packed_ragged_graph) { sdpa_backward_options.set_max_total_seq_len_kv(s_kv); } @@ -1185,12 +1203,10 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i const size_t max_seqlen_q = cfg.max_seqlen_q; 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); void *devPtrQ = input_Q->data.dptr; void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; @@ -1219,13 +1235,14 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); + // These have to match the shape the forward graph declares for Stats and Max, which is + // packed only where the architecture supports it; see derive_f16_fwd_graph_inputs. + const bool use_ragged_stats = + cfg.is_ragged_q && supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - // sm120 does not use ragged stats: the graph declares a dense - // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below). - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if (use_ragged_stats) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1235,8 +1252,7 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if (use_ragged_stats) { output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 01eb4a960c..7f5a461b05 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1124,13 +1124,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_4 = False # Filter: QKV layout - if qkv_format == "thd": + if "thd" in (q_format, kv_format): if pad_between_seqs: if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed): logger.debug( - "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is " + "Disabling FlashAttention 2 and 4 when Q or KV uses THD and there is " "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]" ) use_flash_attention_2 = False @@ -1143,7 +1143,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( - "Disabling FusedAttention as qkv_format = thd is" + "Disabling FusedAttention when Q or KV uses THD because it is" " not supported for compute capability = sm120 and cuDNN version < 9.18.1" ) use_fused_attention = False From 35d6da188d9d97fb6a389af8347cdbba1e62e844 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 14 Aug 2026 08:37:23 -0700 Subject: [PATCH 78/88] address review comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/run_graph_cache.py | 127 +++++++++++++++++ tests/pytorch/attention/test_attention.py | 130 ++++++++++++++++++ .../common/fused_attn/fused_attn.cpp | 44 +++--- .../common/fused_attn/fused_attn_fp8.cu | 24 ++-- .../common/fused_attn/graph_cache.h | 11 +- .../common/fused_attn/graph_cache_debug.h | 29 ++++ 6 files changed, 336 insertions(+), 29 deletions(-) create mode 100644 tests/pytorch/attention/run_graph_cache.py diff --git a/tests/pytorch/attention/run_graph_cache.py b/tests/pytorch/attention/run_graph_cache.py new file mode 100644 index 0000000000..b4c464d3d4 --- /dev/null +++ b/tests/pytorch/attention/run_graph_cache.py @@ -0,0 +1,127 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Worker for test_attention.py::test_fused_attn_graph_cache. + +Runs a fixed sequence of support queries and executions against the cuDNN graph cache and +marks each phase boundary on stderr, so that the parent can attribute the +[FUSED-ATTN-CACHE] lines NVTE_FUSED_ATTN_CACHE_DEBUG=2 emits to the phase that produced +them. + +This runs as its own process because the cache is process-wide and its counters only +accumulate: inside the pytest process, the graphs every earlier test built would be mixed +into the counts, and the cache would already be warm for whatever this test asked about. + +The phases, in order, and what each one is for: + query the first support query for a config -- the miss that builds its graphs + requery the identical query again -- must be answered from the cache + exec forward and backward of that config -- must reuse the graphs the query built, + and is where the plan build the query deferred happens + rescale the same execution with only softmax_scale changed -- must still reuse them, + since attn_scale is normalized out of the cache key + reshape a query differing in max_seqlen -- must build again, once per pass + +Prints ``[CACHE-TEST] fused=1`` (or 0) on stdout so the parent can skip rather than fail on +a GPU or cuDNN version with no fused-attention backend for the config. +""" + +import os +import pathlib +import sys + +import torch + +_current_file = pathlib.Path(__file__).resolve() +sys.path = [str(_current_file.parent.parent)] + sys.path + +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends +from utils import ModelConfig, get_available_attention_backends + +DTYPE = torch.bfloat16 +QKV_FORMAT = "bshd" +QKV_LAYOUT = "bshd_bshd_bshd" + + +def mark_phase(name: str) -> None: + """Delimit the cache events of one phase from the next one's. + + Written to stderr, which is where the diagnostics go, so that the marker keeps its + place in the stream instead of racing them on a second file descriptor. + """ + sys.stderr.write(f"[CACHE-TEST] phase={name}\n") + sys.stderr.flush() + + +def query(config: ModelConfig) -> bool: + """Run one backend support query, as the test suite does, and report whether cuDNN + took the configuration. This is the call that populates the cache without executing + anything.""" + available_backends, _, fused_attn_backends = get_available_attention_backends( + config, qkv_dtype=DTYPE, qkv_layout=QKV_LAYOUT + ) + _, fused_attn_supported, _ = available_backends + return fused_attn_supported and len(fused_attn_backends) > 0 + + +def execute(config: ModelConfig, softmax_scale: float) -> None: + """Run a forward and backward pass of `config` on the fused backend.""" + block = DotProductAttention( + config.num_heads, + config.head_dim_qk, + attention_dropout=config.dropout_p, + qkv_format=QKV_FORMAT, + attn_mask_type=config.attn_mask_type, + softmax_scale=softmax_scale, + layer_number=1, + attention_type=config.attn_type, + ).to(dtype=DTYPE, device="cuda") + shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) + q, k, v = [torch.randn(shape, dtype=DTYPE, device="cuda", requires_grad=True) for _ in range(3)] + out = block(q, k, v, core_attention_bias_type=config.attn_bias_type) + out.backward(torch.randn_like(out)) + # The counters are incremented from the launching thread, but the graphs are not + # necessarily done with; synchronize so that nothing lands in the next phase. + torch.cuda.synchronize() + + +def main() -> int: + torch.manual_seed(1234) + # No mask, no bias, no dropout: the simplest configuration cuDNN supports, so that the + # counts this produces are about the cache rather than about which graph got built. + config = ModelConfig(2, 512, 8, 64) + reshaped = ModelConfig(2, 256, 8, 64) + + mark_phase("query") + fused_available = query(config) + print(f"[CACHE-TEST] fused={int(fused_available)}", flush=True) + if not fused_available: + return 0 + + mark_phase("requery") + query(config) + + # get_available_attention_backends() enables every backend so it can report on all of + # them; the execution phases have to land on the fused one for their cache events to + # exist at all, so leave it the only one available. + 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 + + mark_phase("exec") + execute(config, softmax_scale=0.125) + + mark_phase("rescale") + execute(config, softmax_scale=0.25) + + mark_phase("reshape") + query(reshaped) + + mark_phase("done") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a800df578c..3a6cc17d48 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,9 +1,12 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +import collections import copy import logging import os +import re +import subprocess import sys import pathlib import copy @@ -272,6 +275,133 @@ def test_dpa_checkpoint(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, None, False, False) +# One [FUSED-ATTN-CACHE] event, as either a counter line ("fwd BUILD") or a level-2 trace +# line ("fwd MISS"). The pass and the event name are all this test reads; the trace line's +# cache key is kept so that distinct keys can be counted. +_CACHE_EVENT = re.compile( + r"\[FUSED-ATTN-CACHE\].*\|\s+(?Pfwd|bwd)\s+" + r"(?PBUILD|PLANS|UNSUP|EXEC|MISS|HIT|NOSUP)\b(?P.*)" +) +_CACHE_PHASE = re.compile(r"\[CACHE-TEST\] phase=(?P\w+)") + + +def _parse_cache_events(stderr: str): + """Group the worker's cache diagnostics by the phase that produced them. + + Returns (events, miss_keys): events[phase][(pass, event)] is a count, and + miss_keys[phase][pass] is the set of distinct cache keys that missed, so that "one extra + graph" can be told apart from "the same graph rebuilt". + """ + events = collections.defaultdict(collections.Counter) + miss_keys = collections.defaultdict(lambda: collections.defaultdict(set)) + phase = None + for line in stderr.splitlines(): + phase_match = _CACHE_PHASE.search(line) + if phase_match is not None: + phase = phase_match.group("name") + continue + event_match = _CACHE_EVENT.search(line) + if event_match is None or phase is None: + continue + pass_name, event = event_match.group("pass"), event_match.group("event") + events[phase][(pass_name, event)] += 1 + if event == "MISS": + miss_keys[phase][pass_name].add(event_match.group("rest").split("|")[-1].strip()) + return events, miss_keys + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +def test_fused_attn_graph_cache(): + """Test that the cuDNN graph cache is hit when it should be, and missed when it must be. + + A cuDNN graph build is the most expensive thing in a fused-attention call, so what this + checks is that each distinct configuration pays for one and no more: that a support query + builds the graph an execution then reuses, that a field the graph does not read (here + softmax_scale, which the key normalizes away) does not multiply the cache, and that a + field it does read still gets its own graph. The counters come from + NVTE_FUSED_ATTN_CACHE_DEBUG, which is also what a user would reach for to answer the + same question about their own model. + + The work runs in a subprocess (run_graph_cache.py): the cache is process-wide with + accumulating counters, so within pytest the graphs built by other tests would be + indistinguishable from this test's own. + """ + if torch.cuda.device_count() == 0: + pytest.skip("No CUDA device available.") + + worker = _current_file.parent / "run_graph_cache.py" + result = subprocess.run( + [sys.executable, str(worker)], + env={ + **os.environ, + # Level 2: the per-lookup HIT/MISS lines are what make the hits visible, and + # the volume is trivial for the handful of configurations below. + "NVTE_FUSED_ATTN_CACHE_DEBUG": "2", + "PYTHONUNBUFFERED": "1", + }, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + assert result.returncode == 0, ( + f"{worker.name} failed with exit code {result.returncode}\n" + f"--- stdout ---\n{result.stdout}\n--- stderr (tail) ---\n{result.stderr[-4000:]}" + ) + if "[CACHE-TEST] fused=1" not in result.stdout: + pytest.skip("No cuDNN fused attention backend for the graph cache test config.") + + events, miss_keys = _parse_cache_events(result.stderr) + context = f"\n--- stderr ---\n{result.stderr[-8000:]}" + for phase in ("query", "requery", "exec", "rescale", "reshape"): + assert phase in events, f"worker emitted no cache events for phase {phase}{context}" + + # Both passes are queried by one call and executed by one forward/backward pair, so each + # of them sees the same sequence of events. + for pass_name in ("fwd", "bwd"): + + def count(phase, event, pass_name=pass_name): + return events[phase][(pass_name, event)] + + # The first query builds each pass's graph, and no more than its graph: a support + # query stops at check_support(), leaving the kernel compilation (PLANS) to whoever + # executes it. + assert count("query", "MISS") == 1, f"{pass_name}: expected one cold miss{context}" + assert count("query", "BUILD") == 1, f"{pass_name}: expected one build{context}" + assert count("query", "UNSUP") == 0, f"{pass_name}: cuDNN refused the config{context}" + assert count("query", "PLANS") == 0, f"{pass_name}: query compiled kernels{context}" + + # Asking the identical question again must cost nothing. + assert count("requery", "MISS") == 0, f"{pass_name}: repeated query missed{context}" + assert count("requery", "BUILD") == 0, f"{pass_name}: repeated query rebuilt{context}" + assert count("requery", "HIT") >= 1, f"{pass_name}: repeated query never looked{context}" + + # The execution must find the graph the query left behind -- a miss here is the + # probe/execute key drift this cache is most likely to develop -- and it is what + # finishes the build, exactly once. + assert ( + count("exec", "MISS") == 0 + ), f"{pass_name}: execution missed the query's graph{context}" + assert count("exec", "BUILD") == 0, f"{pass_name}: execution rebuilt the graph{context}" + assert count("exec", "EXEC") >= 1, f"{pass_name}: fused attention never ran{context}" + assert count("exec", "PLANS") == 1, f"{pass_name}: expected one plan build{context}" + + # softmax_scale reaches the graph as a pointer, not as a shape, so the key drops it: + # a different scale has to reuse everything, down to the compiled kernels. + assert count("rescale", "MISS") == 0, f"{pass_name}: attn_scale changed the key{context}" + assert count("rescale", "BUILD") == 0, f"{pass_name}: attn_scale forced a build{context}" + assert count("rescale", "PLANS") == 0, f"{pass_name}: attn_scale recompiled{context}" + assert count("rescale", "EXEC") >= 1, f"{pass_name}: rescaled run did not execute{context}" + + # max_seqlen is a dimension the graph is built at, so it must miss -- once, for one + # new graph, rather than invalidating what is already cached. + assert count("reshape", "MISS") == 1, f"{pass_name}: expected one miss{context}" + assert count("reshape", "BUILD") == 1, f"{pass_name}: expected one build{context}" + assert ( + len(miss_keys["reshape"][pass_name]) == 1 + ), f"{pass_name}: more than one new cache key{context}" + + model_configs_max_logit = { # test: ModelConfig(b, sq, hq, dqk) "max_logit_1": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096), diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index f7dacb9cdc..2d24826d3d 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -184,7 +184,7 @@ NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout) { // map NVTE_QKV_Layout to NVTE_QKV_Format for Q NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); switch (qkv_format) { case NVTE_QKV_Format::NVTE_SBHD: case NVTE_QKV_Format::NVTE_SBHD_2BSHD: @@ -206,7 +206,7 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { // map NVTE_QKV_Layout to NVTE_QKV_Format for KV NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); switch (qkv_format) { case NVTE_QKV_Format::NVTE_SBHD: case NVTE_QKV_Format::NVTE_BSHD_2SBHD: @@ -309,14 +309,16 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); if (is_fp8) { - if (cfg.return_max_logit) { - set_message(message, "FP8 fused attention does not support return_max_logit=True."); + if (set_message_if(cfg.return_max_logit, 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)) + "."); + if (set_message_if(qkv_format != NVTE_QKV_Format::NVTE_BSHD && + qkv_format != NVTE_QKV_Format::NVTE_SBHD && + qkv_format != NVTE_QKV_Format::NVTE_BHSD, + message, + "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(qkv_format)) + ".")) { return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (cfg.check_for_forward_support) { @@ -337,13 +339,21 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi } if (is_f16_or_bf16) { - if (cudnn_runtime_version <= 91500 && cfg.is_training && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - (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."); + // TODO(cyanguwa): re-validate BRCM + cross-attention on sm100 with cuDNN <= 9.7. The + // hand-written support matrix this function replaced rejected bottom-right-diagonal masks + // with max_seqlen_q != max_seqlen_kv there, for a cuDNN bug fixed in 9.7. cuDNN's own + // check_support is the authority now, so the guard is gone; it needs to come back as an + // explicit rejection here, like the CUDA-graph one below, if that bug is a wrong-result + // bug rather than a support gap check_support reports for itself. + if (set_message_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, + message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN.")) { return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } if (cfg.check_for_forward_support) { @@ -455,7 +465,7 @@ 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 { - const char *reject_reason = + const char *const 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"; @@ -582,7 +592,7 @@ 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 { - const char *reject_reason = + const char *const 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"; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index ad18a19481..4214be8614 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -211,9 +211,9 @@ static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& scale_o = mha_graph->tensor(1.0f); } } else if (is_mxfp8) { - NVTE_QKV_Format q_scale_inv_format = + const 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 = + const NVTE_QKV_Format kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.kv_format; std::vector q_scale_strides(4); std::vector k_scale_strides(4); @@ -493,8 +493,9 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + // TODO(cyanguwa): pass bucketed_batch_size cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) + b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -758,13 +759,13 @@ static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& scale_dV = mha_graph->tensor(1.0f); } } else if (is_mxfp8) { - NVTE_QKV_Format q_format = cfg.q_format; - NVTE_QKV_Format kv_format = cfg.kv_format; - NVTE_QKV_Format q_scale_inv_format = + const NVTE_QKV_Format q_format = cfg.q_format; + const NVTE_QKV_Format kv_format = cfg.kv_format; + const 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 = + const NVTE_QKV_Format kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; - NVTE_QKV_Format do_scale_format_ = + const NVTE_QKV_Format do_scale_format_ = (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; // Q_t, K_t, dO_t, dO_f16 std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); @@ -1164,8 +1165,9 @@ void fused_attn_fp8_bwd_impl( const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + // TODO(cyanguwa): pass bucketed_batch_size cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) + b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1276,7 +1278,7 @@ void fused_attn_fp8_fwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const size_t workspace_size = 0; - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + 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)) { fused_attn::fused_attn_fp8_fwd_impl( @@ -1388,7 +1390,7 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const size_t workspace_size = 0; - NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); + const 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( diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index 83a34d2a95..98d8803a14 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -95,11 +95,16 @@ struct CachedGraph { // is what it leaves behind -- cuDNN's own account of the refusal, which is the entire useful // output of a failed query, so nothing is lost by answering from it. Reasons are short strings // and there is one per refused key, so this grows far slower than the graphs beside it. +// Holding the lock and the maps together is also what fixes their relative lifetimes. Members +// are destroyed in reverse declaration order, so the mutex is declared first to be destroyed +// last: the maps go while their guard is still valid, rather than the other way round. Declaring +// a cache and its lock as two separate objects leaves that ordering to whoever writes the next +// one; declaring them here settles it once. template struct GraphCache { + std::mutex mutex; // guards both maps below std::map>> supported; std::map unsupported; - std::mutex mutex; // guards both maps }; // Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: @@ -189,6 +194,10 @@ std::shared_ptr> get_or_build_cached_graph( } else if (refused) { outcome = LookupResult::Unsupported; } + // Recorded after the lock is dropped, so that writing a trace line cannot hold up threads + // querying other keys. The counters are exact, but two lookups that raced on the lock can be + // recorded in the opposite order, so read a level-2 trace as the set of lookups that happened + // rather than as the sequence they happened in. graph_cache_debug::record_cache_lookup(pass, outcome, key); if (cached != nullptr) return cached; diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index a23812dc67..9c4abe1566 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -26,6 +26,35 @@ // An optional ":" suffix picks which processes emit, defaulting to rank 0 // so that output does not scale with the world size: "1:all" for every rank, // "2:0,3" for a specific set. See `rank_selected` for when overriding pays off. +// +// Level 1 on one training step of a supported configuration. Every line begins with +// "[FUSED-ATTN-CACHE] pid=[ rank=] | ", elided below, and carries the running +// totals, of which only the pass being reported is shown (the counters are printed +// right-aligned in a fixed width, dropped here): +// +// THREAD | tid=0 os_tid=1234 +// fwd BUILD | tid=0 | fwd miss=1, hit=0, built=1, unsup=0, plans=0, exec=0 | bwd ... +// bwd BUILD | tid=0 | fwd ... | bwd miss=1, hit=0, built=1, unsup=0, plans=0, exec=0 +// fwd PLANS | tid=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=0 | bwd ... +// ===== summary begin ===== +// SUMMARY-TID | tid=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... +// SUMMARY | tid=all | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... +// fwd check_support | calls=1 | time= 42.135 ms/call +// ===== summary end ===== +// +// Two forward lookups against one build is the shape of a healthy run: the support +// query missed and built, and the execution that followed hit the entry the query left +// behind. `built=1, plans=1` says that graph went on to be executed; `built` above +// `plans` counts graphs built for a query and never run. A refused configuration reads +// `miss=1, unsup=1, built=0` instead, and stays at one refusal however many times it is +// queried. +// +// Level 2 adds one line per lookup and per execution, with the key that decided it: +// +// fwd MISS | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... +// fwd HIT | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... +// +// where diffing two MISS lines names the fields that cost the extra build. // ============================================================================ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ From 24f53ba204189de4ac760707c1e71ae5069f0b19 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Sat, 15 Aug 2026 00:47:33 -0700 Subject: [PATCH 79/88] fix CI failures Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/attention/run_graph_cache.py | 10 +++++++++- transformer_engine/common/fused_attn/fused_attn.cpp | 11 +++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/run_graph_cache.py b/tests/pytorch/attention/run_graph_cache.py index b4c464d3d4..dcf1f1a0a5 100644 --- a/tests/pytorch/attention/run_graph_cache.py +++ b/tests/pytorch/attention/run_graph_cache.py @@ -43,6 +43,14 @@ QKV_FORMAT = "bshd" QKV_LAYOUT = "bshd_bshd_bshd" +# Derived exactly as DotProductAttention derives it, because the query below has to ask about +# the configuration the execution phases will run: deterministic is part of the backward cache +# key, so a query that assumed the default would build a graph the execution then misses. +DETERMINISTIC = ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() +) + def mark_phase(name: str) -> None: """Delimit the cache events of one phase from the next one's. @@ -59,7 +67,7 @@ def query(config: ModelConfig) -> bool: took the configuration. This is the call that populates the cache without executing anything.""" available_backends, _, fused_attn_backends = get_available_attention_backends( - config, qkv_dtype=DTYPE, qkv_layout=QKV_LAYOUT + config, qkv_dtype=DTYPE, qkv_layout=QKV_LAYOUT, deterministic=DETERMINISTIC ) _, fused_attn_supported, _ = available_backends return fused_attn_supported and len(fused_attn_backends) > 0 diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 2d24826d3d..e808305546 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -297,6 +297,17 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } + // Ragged Q/KV requires sm90+. cuDNN's check_support accepts them below that, but the only + // graph we can build there is the dense max_seqlen one -- supports_packed_ragged_graph() is + // false, so SDPA_backward never gets max_total_seq_len_q/kv and its dQ/dK/dV come back wrong. + // A wrong-result rejection like this one has to be stated here; check_support answers whether + // cuDNN can run the graph, not whether the graph computes what we asked for. + if (set_message_if((cfg.is_ragged_q || cfg.is_ragged_kv) && + cuda::sm_arch(cuda::current_device()) < 90, + message, "Ragged (THD) Q or KV requires compute capability 9.0 or higher.")) { + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + // TE's cuDNN fused-attention graph does not represent pre-scale bias. if (set_message_if(cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS, message, "Fused attention does not support pre-scale bias.")) { From 2a7eed9eb5faa0fc174b63103e018d23fb7b24ed Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 17 Aug 2026 07:41:08 -0700 Subject: [PATCH 80/88] more fixes: cache only genuine cuDNN refusals, bound the graph cache with LRU, drop cudnn.h from the public header, guard move-assignments against NULL Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 22 ++- docs/examples/attention/attention.ipynb | 16 +- .../common/fused_attn/config_and_params.cpp | 6 + .../common/fused_attn/config_and_params.h | 21 ++ .../common/fused_attn/fused_attn.cpp | 24 ++- .../common/fused_attn/fused_attn_fp8.h | 2 + .../common/fused_attn/graph_cache.h | 181 +++++++++++++++--- .../common/fused_attn/graph_cache_debug.h | 89 ++++++--- .../include/transformer_engine/fused_attn.h | 22 ++- 9 files changed, 311 insertions(+), 72 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index f1054084e4..0cb4288f81 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -198,9 +198,27 @@ backend-selection overview. .. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG - :Type: ``int`` (0 or 1) + :Type: ``int`` (0, 1 or 2), optionally followed by ``:`` :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. + :Description: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). Output goes to stderr, prefixed ``[FUSED-ATTN-CACHE]``. + + ``1`` emits one line per event that happens once per distinct cache key -- ``BUILD`` when a graph is constructed, ``PLANS`` when its kernels are compiled on first execution, ``UNSUP`` when cuDNN refuses a configuration -- plus an end-of-run ``SUMMARY`` (aggregate and per thread) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. + + ``2`` additionally emits a per-lookup ``HIT``/``MISS``/``NOSUP`` line carrying the full cache key, and a per-execution ``EXEC`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. + + By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. + + Has negligible overhead when unset. + +.. envvar:: NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES + + :Type: ``int`` + :Default: ``500`` + :Description: Ceiling on the number of entries the FusedAttention graph cache keeps per build site, evicting least-recently-used entries to stay under it. There are four build sites (F16 and FP8, forward and backward), each holding a cache of graphs and a cache of configurations cuDNN refused, and the ceiling applies to each of those independently. + + The default is meant to be out of the way of real work rather than tight: a training step reuses a handful of configurations and an inference server with bucketed sequence lengths tens of them, so a few hundred is already more shape diversity than a model exhibits. The ceiling exists for workloads whose key space is effectively unbounded -- a test suite sweeping shapes, or a server keying on something that never repeats -- where an unbounded cache is a slow leak of cuDNN graphs and their execution plans for the life of the process. + + Set to ``0`` to disable the ceiling entirely, for a workload that genuinely has thousands of live configurations and would rather spend the memory than rebuild. Evicting a graph never disturbs one that is executing; execution holds its own reference. .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 4ffa804401..79ba646ca6 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -253,7 +253,21 @@ "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", - "" + "\n", + "\n", + "Once cuDNN attention has been selected, a separate flag reports on the cuDNN graph cache underneath it. Transformer Engine builds a cuDNN graph per distinct attention configuration and reuses it, so a workload that unexpectedly rebuilds graphs pays for it on every step.\n", + "```\n", + "NVTE_FUSED_ATTN_CACHE_DEBUG = 0/1/2 # disables/enables graph cache diagnostics\n", + "```\n", + "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]`, and there is one per event that happens once per configuration: `BUILD` when a graph is constructed, `PLANS` when its kernels are compiled on first execution, and `UNSUP` when cuDNN declines a configuration. An end-of-run `SUMMARY` gives the totals, per thread and per device, followed by where the build time went:\n", + "```\n", + "[FUSED-ATTN-CACHE] pid=1234 | fwd BUILD | tid=0 dev=0 | fwd miss=1, hit=0, built=1, ...\n", + "[FUSED-ATTN-CACHE] pid=1234 | SUMMARY | tid=all dev=all | fwd miss=2, hit=1, built=1, ...\n", + "[FUSED-ATTN-CACHE] pid=1234 | fwd check_support | calls=1 | time= 42.135 ms/call\n", + "```\n", + "The number to read first is `built`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n", + "\n", + "This flag is supported in both PyTorch and JAX, since the cache it reports on lives in the common C++ layer. By default only rank 0 emits; see [NVTE_FUSED_ATTN_CACHE_DEBUG](../../envvars.rst) for selecting other ranks, and for the companion `NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES` which bounds how many configurations the cache retains." ] }, { diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index f42850ddad..8a07c14121 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -138,6 +138,12 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { // attn_scale is a pass-by-value graph input and different scales can share the same cached graph cache_cfg.attn_scale = 1.0f; + // cuda_graph never reaches a graph builder. Its one use is the cuDNN <= 9.15 rejection in + // nvte_get_fused_attn_backend_v2(), which is decided before the cache is consulted, so a + // configuration that gets this far builds the same graph either way. Left in the key it would + // give a workload that both captures and runs eagerly two entries for every configuration. + cache_cfg.cuda_graph = false; + // Restrict each direction's key to the fields its graph actually consumes, so // no redundant graphs are built and no cache misses either if (check_for_forward_support) { diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 3bbf77dfa3..d116893adf 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -165,6 +165,15 @@ struct FusedAttnConfig { sizeof(size_t), // bias_seqlen_kv }; + // The public header asks contributors to append to NVTEFusedAttnConfigAttribute, and the + // accessors index attr_sizes[attr] after checking only that attr is below the sentinel. An + // enumerator added without its size here would therefore read one past the end of this array, + // silently and only for the new attribute. Tying the two together turns that into a build + // failure at the line that has to change. + static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnConfigNumAttributes, + "attr_sizes must have one entry per NVTEFusedAttnConfigAttribute; add the size of " + "the new attribute alongside its enumerator."); + 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, @@ -302,6 +311,12 @@ struct FusedAttnFwdParams { sizeof(cudaStream_t), // stream }; + // See FusedAttnConfig::attr_sizes: an enumerator appended without a size here reads past the + // end of this array. + static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnFwdParamsNumAttributes, + "attr_sizes must have one entry per NVTEFusedAttnFwdParamsAttribute; add the size " + "of the new attribute alongside its enumerator."); + // 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 @@ -398,6 +413,12 @@ struct FusedAttnBwdParams { sizeof(cudaStream_t), // stream }; + // See FusedAttnConfig::attr_sizes: an enumerator appended without a size here reads past the + // end of this array. + static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnBwdParamsNumAttributes, + "attr_sizes must have one entry per NVTEFusedAttnBwdParamsAttribute; add the size " + "of the new attribute alongside its enumerator."); + // 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 diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index e808305546..c47dc9b603 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -254,6 +254,7 @@ bool set_message_if(bool rejected, const char **message, std::string reason) { // 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) { + NVTE_API_CALL(nvte_get_fused_attn_backend_v2); using namespace transformer_engine; using namespace transformer_engine::fused_attn; // Derived on a copy, leaving the caller's config untouched: this function answers a question @@ -297,11 +298,24 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - // Ragged Q/KV requires sm90+. cuDNN's check_support accepts them below that, but the only - // graph we can build there is the dense max_seqlen one -- supports_packed_ragged_graph() is - // false, so SDPA_backward never gets max_total_seq_len_q/kv and its dQ/dK/dV come back wrong. - // A wrong-result rejection like this one has to be stated here; check_support answers whether - // cuDNN can run the graph, not whether the graph computes what we asked for. + // Ragged Q/KV requires sm90+, the rule the hand-written support matrix this function replaced + // carried as `qkv_format == NVTE_THD && sm_arch_ >= 90`. Below sm90 the only graph we can build + // is the dense max_seqlen one -- supports_packed_ragged_graph() is false -- so SDPA_backward + // never gets max_total_seq_len_q/kv and its dQ/dK/dV come back wrong. + // + // This is ours to state because it is a wrong-result rejection, and check_support answers a + // different question: whether cuDNN can run the graph, not whether the graph computes what we + // asked for. cuDNN's own answer has moved, which is what makes the distinction worth spelling + // out here. Its frontend gates ragged SDPA on `sm < 90 && cudnn < 9.18.1`, so through 9.18.0 it + // would have refused this configuration for us and the rule below is redundant; from 9.18.1 it + // accepts sm80/sm89 ragged and the rule is the only thing standing between a THD model on an + // A100 and silently wrong gradients. Lifting it is a change to the graphs TE builds -- packed + // ragged shapes, and the Stats/LSE layouts that go with them, which cuDNN documents as + // differing on sm8x -- not a change to this condition, and that work is deliberately not part + // of this refactor. + // + // sm120 takes that same dense path and is left enabled, as it was before this refactor; + // whether it has the same problem is a separate question from restoring the sm90 rule. if (set_message_if((cfg.is_ragged_q || cfg.is_ragged_kv) && cuda::sm_arch(cuda::current_device()) < 90, message, "Ragged (THD) Q or KV requires compute capability 9.0 or higher.")) { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 79b279a833..9e8f997d98 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -11,6 +11,8 @@ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ +#include + #include #include "config_and_params.h" diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index 98d8803a14..dd7f78b11c 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -22,6 +22,8 @@ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ +#include +#include #include #include #include @@ -44,13 +46,46 @@ namespace fused_attn { // configuration and reproducible for a given key, so it can be remembered and replayed, whereas // a failure that came from the machine's state at that moment (an allocation that did not fit, a // CUDA error left behind by unrelated work) could well succeed on the next attempt and must not -// be turned into a permanent answer. Only the four adjudicating frontend calls in -// validate_and_check_support() raise this; every other failure keeps its ordinary type and is +// be turned into a permanent answer. Only a frontend call that returned one of the codes +// is_unsupported_verdict() names raises this; every other failure keeps its ordinary type and is // re-attempted the next time the key comes around. struct UnsupportedGraph : public std::runtime_error { explicit UnsupportedGraph(const std::string &reason) : std::runtime_error(reason) {} }; +// Whether a frontend error code is a verdict on the graph rather than a report of something +// that went wrong on the way to reaching one. +// +// cudnn-frontend distinguishes the two, and the negative cache is only sound for the first. +// Three codes are verdicts: +// +// GRAPH_NOT_SUPPORTED is what the frontend's own support surface returns, from validate(). +// Nearly every rule it checks by hand -- the architecture gates, the head-dim limits, the +// version-specific workarounds -- reports itself this way. +// GRAPH_EXECUTION_PLAN_CREATION_FAILED is what check_support() returns when no engine config +// cuDNN's heuristics offered can run the graph. This is the verdict for everything the +// frontend does not rule on itself and defers to the backend, so leaving it out would +// exclude most of what a support probe actually discovers. +// UNSUPPORTED_GRAPH_FORMAT is a verdict by name and costs nothing to accept, though no +// frontend release we build against returns it. +// +// All three are properties of the key and will be just as true the next time it is asked. Every +// other code -- CUDNN_BACKEND_API_FAILED, CUDA_API_FAILED, HEURISTIC_QUERY_FAILED, HANDLE_ERROR, +// INVALID_CUDA_DEVICE and the rest -- describes the process at that moment: an OOM under memory +// pressure, a sticky CUDA error left by unrelated work, a handle on the wrong device. +// CUDNN_BACKEND_API_FAILED is the one to be careful about, since the frontend raises it for any +// non-success cudnnStatus_t and so cannot tell CUDNN_STATUS_ALLOC_FAILED from +// CUDNN_STATUS_NOT_SUPPORTED; caching it would let a moment of memory pressure blacklist a +// configuration that is genuinely supported, for the life of the process, and the wider the +// cache's reach the worse that gets. So those are raised as ordinary errors, which leave nothing +// behind and are retried when the key next comes around. Either way cuDNN's own message reaches +// the caller; only whether it is remembered differs. +inline bool is_unsupported_verdict(cudnn_frontend::error_code_t code) { + return code == cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED || + code == cudnn_frontend::error_code_t::GRAPH_EXECUTION_PLAN_CREATION_FAILED || + code == cudnn_frontend::error_code_t::UNSUPPORTED_GRAPH_FORMAT; +} + // The reason string an is_supported_* helper reports for `e`: its message, or `fallback` if it // has none. Those helpers signal support by returning the empty string, so a refusal that // arrives without an explanation would be read as an endorsement and the caller would go on to @@ -86,7 +121,10 @@ struct CachedGraph { // One build site's cache. Process-wide rather than per-thread so that a graph is reused // across threads instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a -// shared plan, and cudnn-frontend >= 1.25.0 has a thread-safe execute(). +// shared plan, and the frontend's execute() builds its variant pack in a local rather than in +// the graph, so it does not write to the shared object. No particular frontend version is +// relied on for that -- it has held for far longer than the >= 1.25.0 the build requirements +// ask for, which is there for unrelated features. // // Refusals are cached alongside the graphs, under the same keys and the same lock. A support // query for an unsupported configuration is otherwise the most expensive thing this cache sees: @@ -100,13 +138,73 @@ struct CachedGraph { // last: the maps go while their guard is still valid, rather than the other way round. Declaring // a cache and its lock as two separate objects leaves that ordering to whoever writes the next // one; declaring them here settles it once. +// +// Both maps are bounded; see cache_capacity(). `last_used` is what makes the bound an LRU rather +// than an arbitrary cull: it is stamped from `clock` on every insertion and every hit, so the +// entry with the smallest value is the one that has gone longest without being asked for. The +// clock is an ordinary member rather than an atomic because it is only ever touched under +// `mutex`, alongside the maps it orders. template struct GraphCache { - std::mutex mutex; // guards both maps below - std::map>> supported; - std::map unsupported; + struct Slot { + std::shared_ptr> entry; + uint64_t last_used; + }; + struct Refusal { + std::string reason; + uint64_t last_used; + }; + + std::mutex mutex; // guards everything below + uint64_t clock = 0; + std::map supported; + std::map unsupported; }; +// The default ceiling on entries in one of the maps of one build site's cache. +// +// Sized to be out of the way of real work rather than to be tight. A training step reuses a +// handful of configurations, an inference server with bucketed sequence lengths tens of them; +// a few hundred is already far more shape diversity than a model exhibits. What the ceiling is +// for is the case where the key space is effectively unbounded -- a test suite sweeping shapes, +// or a serving workload that keys on something that never repeats -- where an unbounded cache +// is a slow leak of cuDNN graphs and their execution plans for the life of the process. +constexpr size_t kDefaultCacheCapacity = 500; + +// The ceiling in force, from NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES if it is set. 0 means no ceiling, +// which is the escape hatch for a workload that genuinely has thousands of live configurations +// and would rather spend the memory than rebuild. Read once: the limit is a property of the run. +inline size_t cache_capacity() { + static const size_t capacity = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES"); + if (e == nullptr || e[0] == '\0') return kDefaultCacheCapacity; + const long long v = std::atoll(e); // NOLINT(runtime/int) + return v < 0 ? kDefaultCacheCapacity : static_cast(v); + }(); + return capacity; +} + +// Make room in `entries` for one more, by dropping the least recently used until there is. +// Call under the cache's lock. +// +// Evicting a graph does not invalidate one that is in use. get_or_build_cached_graph() hands +// back a shared_ptr, so a thread that is executing an entry holds it alive regardless of what +// the map does; erasing here drops the cache's reference and nothing else. The scan is linear, +// but it runs only when the cache is full, and comparing a few hundred integers is nothing +// beside the graph build it is making room for. +template +void evict_to_fit(Map &entries) { + const size_t capacity = cache_capacity(); + if (capacity == 0) return; + while (entries.size() >= capacity) { + auto oldest = entries.begin(); + for (auto it = entries.begin(); it != entries.end(); ++it) { + if (it->second.last_used < oldest->second.last_used) oldest = it; + } + entries.erase(oldest); + } +} + // Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: // validate, build_operation_graph, create_execution_plans, check_support. The sequence is // identical for both passes and both backends, so it is defined once here; `pass` only selects @@ -119,30 +217,40 @@ struct GraphCache { // the graph want the throw as well, since there is nothing useful to do with an unsupported // graph but fail. // -// The throw is re-raised as UnsupportedGraph, which is what marks it cacheable. These four calls -// are cuDNN adjudicating a graph it has been handed, so a failure among them is a statement about -// the graph rather than about the moment -- which is the property the negative cache needs, and -// the reason the boundary is drawn here rather than around a wider region. build_plans() and -// execute() sit outside it: they commit real resources and can fail for reasons that have nothing -// to do with the configuration. +// A failure is raised as UnsupportedGraph only when the frontend's own error code says the graph +// was adjudicated and refused; see is_unsupported_verdict(). Anything else these calls can report +// is a failure to reach a verdict and is raised through NVTE_ERROR, so it is not remembered. +// Classifying on the code rather than on which call failed is what keeps that honest: all four of +// these calls can fail for environmental reasons too -- build_operation_graph() and +// create_execution_plans() both talk to the cuDNN backend -- so their position in the sequence +// says nothing about whether the failure was about the configuration. // -// build_plans() is left out for a second reason as well: it belongs to whoever executes the graph, -// once, the first time it is needed. See CachedGraph. +// build_plans() and execute() sit outside this function entirely: they commit real resources, and +// build_plans() belongs to whoever executes the graph, once, the first time it is needed. See +// CachedGraph. inline void validate_and_check_support(const char *pass, cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { - try { - graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::Validate, - [&] { NVTE_CHECK_CUDNN_FE(graph.validate()); }); - graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::BuildOpGraph, - [&] { NVTE_CHECK_CUDNN_FE(graph.build_operation_graph(handle)); }); - graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::CreatePlans, [&] { - NVTE_CHECK_CUDNN_FE(graph.create_execution_plans({cudnn_frontend::HeurMode_t::A})); - }); - graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::CheckSupport, - [&] { NVTE_CHECK_CUDNN_FE(graph.check_support()); }); - } catch (const std::exception &e) { - throw UnsupportedGraph(e.what()); - } + auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { + cudnn_frontend::error_t error; + graph_cache_debug::timer(pass, stage, [&] { error = call(); }); + if (error.is_good()) return; + // cuDNN normally explains itself; fall back to the call's name so that a refusal can never + // arrive as an empty string, which the is_supported_* helpers would read as an endorsement. + const std::string reason = + error.err_msg.empty() ? std::string(call_name) + " failed." : error.err_msg; + if (is_unsupported_verdict(error.code)) throw UnsupportedGraph(reason); + NVTE_ERROR("cuDNN Error in ", call_name, ": ", reason, + " For more information, enable cuDNN error logging by setting CUDNN_LOGERR_DBG=1 " + "and CUDNN_LOGDEST_DBG=stderr in the environment."); + }; + + run(graph_cache_debug::BuildStage::Validate, "validate", [&] { return graph.validate(); }); + run(graph_cache_debug::BuildStage::BuildOpGraph, "build_operation_graph", + [&] { return graph.build_operation_graph(handle); }); + run(graph_cache_debug::BuildStage::CreatePlans, "create_execution_plans", + [&] { return graph.create_execution_plans({cudnn_frontend::HeurMode_t::A}); }); + run(graph_cache_debug::BuildStage::CheckSupport, "check_support", + [&] { return graph.check_support(); }); } // The cached entry for `key`, building and inserting it via `build` if absent. Throws @@ -172,6 +280,8 @@ std::shared_ptr> get_or_build_cached_graph( GraphCache &cache, const FusedAttnConfig &key, const char *pass, cudnnHandle_t handle, BuildFn &&build) { using Entry = CachedGraph; + using Slot = typename GraphCache::Slot; + using Refusal = typename GraphCache::Refusal; std::shared_ptr cached; bool refused = false; @@ -180,11 +290,15 @@ std::shared_ptr> get_or_build_cached_graph( std::lock_guard lock(cache.mutex); auto it = cache.supported.find(key); if (it != cache.supported.end()) { - cached = it->second; + it->second.last_used = ++cache.clock; + cached = it->second.entry; } else { auto refusal = cache.unsupported.find(key); refused = (refusal != cache.unsupported.end()); - if (refused) reason = refusal->second; + if (refused) { + refusal->second.last_used = ++cache.clock; + reason = refusal->second.reason; + } } } using graph_cache_debug::LookupResult; @@ -216,7 +330,8 @@ std::shared_ptr> get_or_build_cached_graph( } catch (const UnsupportedGraph &e) { { std::lock_guard lock(cache.mutex); - cache.unsupported.insert({key, e.what()}); + evict_to_fit(cache.unsupported); + cache.unsupported.insert({key, Refusal{e.what(), ++cache.clock}}); } graph_cache_debug::record_unsupported(pass); throw; @@ -224,7 +339,11 @@ std::shared_ptr> get_or_build_cached_graph( graph_cache_debug::record_build(pass); { std::lock_guard lock(cache.mutex); - return cache.supported.insert({key, std::move(entry)}).first->second; + evict_to_fit(cache.supported); + // On a losing race the insert does nothing: the temporary Slot is destroyed with the graph + // this thread built, and what comes back is the winner's entry. + auto inserted = cache.supported.insert({key, Slot{std::move(entry), ++cache.clock}}); + return inserted.first->second.entry; } } diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 9c4abe1566..84c51e3cd2 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -32,16 +32,20 @@ // totals, of which only the pass being reported is shown (the counters are printed // right-aligned in a fixed width, dropped here): // -// THREAD | tid=0 os_tid=1234 -// fwd BUILD | tid=0 | fwd miss=1, hit=0, built=1, unsup=0, plans=0, exec=0 | bwd ... -// bwd BUILD | tid=0 | fwd ... | bwd miss=1, hit=0, built=1, unsup=0, plans=0, exec=0 -// fwd PLANS | tid=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=0 | bwd ... +// THREAD | tid=0 dev=0 os_tid=1234 +// fwd BUILD | tid=0 dev=0 | fwd miss=1, hit=0, built=1, unsup=0, plans=0, exec=0 | bwd ... +// bwd BUILD | tid=0 dev=0 | fwd ... | bwd miss=1, hit=0, built=1, unsup=0, plans=0, ... +// fwd PLANS | tid=0 dev=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=0 | bwd ... // ===== summary begin ===== -// SUMMARY-TID | tid=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... -// SUMMARY | tid=all | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... +// SUMMARY-TID | tid=0 dev=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... +// SUMMARY | tid=all dev=all | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... // fwd check_support | calls=1 | time= 42.135 ms/call // ===== summary end ===== // +// The device column matters as soon as one process drives more than one -- device_id is part of +// the cache key, so the same shape on two devices is two entries, and a build count that looks +// doubled is explained by reading which device each BUILD came from. +// // Two forward lookups against one build is the shape of a healthy run: the support // query missed and built, and the execution that followed hit the entry the query left // behind. `built=1, plans=1` says that graph went on to be executed; `built` above @@ -77,6 +81,7 @@ #include #include +#include "../util/cuda_runtime.h" #include "config_and_params.h" namespace transformer_engine { @@ -208,8 +213,15 @@ inline EventCounters &counters(bool is_fwd) { // Per-thread counters, so the summary can break down build/exec/hit/miss by // thread. In the single-process context-parallel case each device is driven by // its own thread, so this reveals which thread built/executed what. +// +// `device` is the device this thread last drove, restamped on every event. The event lines print +// the live current device, which is exact; this exists for the SUMMARY-TID rows, which are +// written at exit by whichever thread is exiting and so cannot ask the recorded thread what it +// was working on. A thread that stays on one device -- which is the arrangement everything here +// is built around, device_id being part of the cache key -- makes the two the same answer. struct ThreadCounters { unsigned tid = 0; + std::atomic device{-1}; EventCounters fwd; EventCounters bwd; }; @@ -239,13 +251,19 @@ inline ThreadCounters &thread_counters() { static thread_local ThreadCounters *tc = [] { auto *p = new ThreadCounters(); p->tid = thread_seq_id(); + // Stamped here as well as on every event, so that a thread which only ever hits the cache -- + // and so never reaches print_counters() at level 1 -- still names a device in the summary + // rather than reporting the -1 it was constructed with. + const int device = cuda::current_device(); + p->device.store(device, std::memory_order_relaxed); { std::lock_guard lock(thread_registry_mutex()); thread_registry().push_back(p); } // One line per thread, mapping the short id to something nsys/gdb can match. - std::fprintf(stderr, "[FUSED-ATTN-CACHE] %s | THREAD | tid=%-3u os_tid=%" PRId64 "\n", - process_tag().c_str(), p->tid, os_thread_id()); + std::fprintf(stderr, + "[FUSED-ATTN-CACHE] %s | THREAD | tid=%-3u dev=%-3d os_tid=%" PRId64 "\n", + process_tag().c_str(), p->tid, device, os_thread_id()); std::fflush(stderr); return p; }(); @@ -259,41 +277,52 @@ inline EventCounters &thread_counters(bool is_fwd) { // Format one counter block (aggregate or a single thread's) as one line. // `tid_field` is the whole thread column, e.g. "tid=3"; the aggregate row passes -// "tid=all" so that it cannot be misread as thread 0's row. +// "tid=all" so that it cannot be misread as thread 0's row. `dev_field` is the device column and +// works the same way, "dev=all" on the aggregate row -- the counters there are summed across +// whatever devices the process drove, so naming one of them would be a lie. // // The columns are meant to be read against two identities. Every lookup lands in exactly one of // miss and hit, and every miss ends in exactly one of built and unsup -- so `miss = built + unsup` // and a shortfall in either means a build died of something other than a refusal. `built >= plans` // always, the difference being graphs that a support query built and nothing has yet run. inline std::string format_counter_line(const char *event, const char *tid_field, - const EventCounters &f, const EventCounters &b) { + const char *dev_field, const EventCounters &f, + const EventCounters &b) { char buf[768]; std::snprintf(buf, sizeof(buf), - "[FUSED-ATTN-CACHE] %s | %-11s | %-7s | fwd miss=%4" PRIu64 ", hit=%4" PRIu64 + "[FUSED-ATTN-CACHE] %s | %-11s | %-7s %-7s | fwd miss=%4" PRIu64 ", hit=%4" PRIu64 ", built=%4" PRIu64 ", unsup=%4" PRIu64 ", plans=%4" PRIu64 ", exec=%4" PRIu64 " | bwd miss=%4" PRIu64 ", hit=%4" PRIu64 ", built=%4" PRIu64 ", unsup=%4" PRIu64 ", plans=%4" PRIu64 ", exec=%4" PRIu64 "\n", - process_tag().c_str(), event, tid_field, f.miss.load(std::memory_order_relaxed), - f.hit.load(std::memory_order_relaxed), f.built.load(std::memory_order_relaxed), - f.unsup.load(std::memory_order_relaxed), f.plans.load(std::memory_order_relaxed), - f.exec.load(std::memory_order_relaxed), b.miss.load(std::memory_order_relaxed), - b.hit.load(std::memory_order_relaxed), b.built.load(std::memory_order_relaxed), - b.unsup.load(std::memory_order_relaxed), b.plans.load(std::memory_order_relaxed), - b.exec.load(std::memory_order_relaxed)); + process_tag().c_str(), event, tid_field, dev_field, + f.miss.load(std::memory_order_relaxed), f.hit.load(std::memory_order_relaxed), + f.built.load(std::memory_order_relaxed), f.unsup.load(std::memory_order_relaxed), + f.plans.load(std::memory_order_relaxed), f.exec.load(std::memory_order_relaxed), + b.miss.load(std::memory_order_relaxed), b.hit.load(std::memory_order_relaxed), + b.built.load(std::memory_order_relaxed), b.unsup.load(std::memory_order_relaxed), + b.plans.load(std::memory_order_relaxed), b.exec.load(std::memory_order_relaxed)); return std::string(buf); } -inline void print_counter_block(const char *event, const char *tid_field, const EventCounters &f, - const EventCounters &b) { - const std::string line = format_counter_line(event, tid_field, f, b); +inline void print_counter_block(const char *event, const char *tid_field, const char *dev_field, + const EventCounters &f, const EventCounters &b) { + const std::string line = format_counter_line(event, tid_field, dev_field, f, b); std::fputs(line.c_str(), stderr); std::fflush(stderr); } +// One event line, from the thread the event happened on. The device is read live rather than +// remembered, so it is the device this event was actually issued against, and is recorded on the +// thread's block on the way past for the benefit of the exit summary. inline void print_counters(const char *event) { + const int device = cuda::current_device(); + thread_counters().device.store(device, std::memory_order_relaxed); char tid_field[16]; + char dev_field[16]; std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); - print_counter_block(event, tid_field, counters(/*is_fwd=*/true), counters(/*is_fwd=*/false)); + std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); + print_counter_block(event, tid_field, dev_field, counters(/*is_fwd=*/true), + counters(/*is_fwd=*/false)); } // A graph built through check_support() and cached. Call after the build, from the miss @@ -467,10 +496,11 @@ inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { // Times one stage: clock read in the constructor, accumulated in the destructor. // Recording on scope exit rather than at an explicit stop() keeps a failing stage -// measurable -- the frontend calls are wrapped in NVTE_CHECK_CUDNN_FE, which -// throws, and the destructor still runs during unwinding -- so a build that dies -// in `check_support` contributes its time to failure instead of vanishing from the -// summary. `on` is latched at construction rather than re-tested in the destructor, +// measurable: `build_plans` throws through NVTE_CHECK_CUDNN_FE and the destructor +// still runs during unwinding, so a build that dies there contributes its time to +// the failure instead of vanishing from the summary. The four stages before it +// return their status instead of throwing, and are timed the same way for the same +// reason. `on` is latched at construction rather than re-tested in the destructor, // which is what keeps that symmetric: the destructor can never accumulate against a // `start` the constructor left unset. struct ScopedBuildTimer { @@ -527,12 +557,15 @@ inline void register_summary_once() { [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); for (const ThreadCounters *tc : blocks) { char tid_field[16]; + char dev_field[16]; std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); - block += format_counter_line("SUMMARY-TID", tid_field, tc->fwd, tc->bwd); + std::snprintf(dev_field, sizeof(dev_field), "dev=%d", + tc->device.load(std::memory_order_relaxed)); + block += format_counter_line("SUMMARY-TID", tid_field, dev_field, tc->fwd, tc->bwd); } } // Totals last, so they read as the sum of the per-thread lines above. - block += format_counter_line("SUMMARY", "tid=all", counters(/*is_fwd=*/true), + block += format_counter_line("SUMMARY", "tid=all", "dev=all", counters(/*is_fwd=*/true), counters(/*is_fwd=*/false)); for (int p = 0; p < 2; ++p) { const bool is_fwd = (p == 0); diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 938fa1747e..d2307b07b6 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -11,8 +11,6 @@ #ifndef TRANSFORMER_ENGINE_FUSED_ATTN_H_ #define TRANSFORMER_ENGINE_FUSED_ATTN_H_ -#include - #include "stdint.h" #include "transformer_engine.h" @@ -1004,7 +1002,13 @@ class FusedAttnConfigWrapper { FusedAttnConfigWrapper &operator=(FusedAttnConfigWrapper &&other) noexcept { if (this != &other) { - nvte_destroy_fused_attn_config(cfg_); + // Guarded as the destructor is. A moved-from wrapper holds nullptr, and the C API rejects a + // NULL handle by throwing; thrown out of a noexcept function that is a call to + // std::terminate, which no caller can catch. The guard belongs on this side rather than in + // nvte_destroy_*, so that the C entry point keeps reporting a genuinely bad handle. + if (cfg_ != nullptr) { + nvte_destroy_fused_attn_config(cfg_); + } cfg_ = other.cfg_; other.cfg_ = nullptr; } @@ -1179,7 +1183,11 @@ class FusedAttnFwdParamsWrapper { FusedAttnFwdParamsWrapper &operator=(FusedAttnFwdParamsWrapper &&other) noexcept { if (this != &other) { - nvte_destroy_fused_attn_fwd_params(params_); + // See FusedAttnConfigWrapper::operator=: destroying a moved-from (NULL) handle throws out + // of a noexcept function, which is std::terminate. + if (params_ != nullptr) { + nvte_destroy_fused_attn_fwd_params(params_); + } params_ = other.params_; other.params_ = nullptr; } @@ -1327,7 +1335,11 @@ class FusedAttnBwdParamsWrapper { FusedAttnBwdParamsWrapper &operator=(FusedAttnBwdParamsWrapper &&other) noexcept { if (this != &other) { - nvte_destroy_fused_attn_bwd_params(params_); + // See FusedAttnConfigWrapper::operator=: destroying a moved-from (NULL) handle throws out + // of a noexcept function, which is std::terminate. + if (params_ != nullptr) { + nvte_destroy_fused_attn_bwd_params(params_); + } params_ = other.params_; other.params_ = nullptr; } From 76ae514891b40b7b555d49c5b509082f39411dbf Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 19 Aug 2026 05:14:37 -0700 Subject: [PATCH 81/88] WIP: restructuring and polishing Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 18 +- docs/examples/attention/attention.ipynb | 12 +- tests/pytorch/attention/test_attention.py | 42 +- .../common/fused_attn/config_and_params.cpp | 80 ++- .../common/fused_attn/config_and_params.h | 50 +- .../common/fused_attn/fused_attn.cpp | 157 +++-- .../fused_attn_f16_arbitrary_seqlen.cu | 243 +++---- .../common/fused_attn/fused_attn_fp8.cu | 126 ++-- .../common/fused_attn/graph_cache.h | 150 ++-- .../common/fused_attn/graph_cache_debug.h | 649 ++++++++++++------ transformer_engine/common/fused_attn/utils.h | 17 +- .../include/transformer_engine/fused_attn.h | 14 +- 12 files changed, 900 insertions(+), 658 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 0cb4288f81..af06ebbb58 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -202,23 +202,17 @@ backend-selection overview. :Default: ``0`` :Description: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). Output goes to stderr, prefixed ``[FUSED-ATTN-CACHE]``. - ``1`` emits one line per event that happens once per distinct cache key -- ``BUILD`` when a graph is constructed, ``PLANS`` when its kernels are compiled on first execution, ``UNSUP`` when cuDNN refuses a configuration -- plus an end-of-run ``SUMMARY`` (aggregate and per thread) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. + ``1`` emits one line per event that happens once per distinct cache key -- ``BUILD_GRAPH`` when a graph is constructed, ``BUILD_PLANS`` when its kernels are compiled on first execution, ``UNSUPPORTED`` when cuDNN refuses a configuration -- plus an end-of-run ``SUMMARY`` (per backend, per thread, and across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries. - ``2`` additionally emits a per-lookup ``HIT``/``MISS``/``NOSUP`` line carrying the full cache key, and a per-execution ``EXEC`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. + Every line names the build site behind it -- ``f16`` or ``fp8``, then the pass -- and carries only that backend's counters, so a process that uses both can still tell which of them built what. A backend the run never reached is left out of the summary entirely. - By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. - - Has negligible overhead when unset. + The two hit columns name which of the cache's two maps answered a lookup: ``hit_supported`` is a cached graph reused, while ``hit_unsupported`` is a configuration cuDNN already refused, replayed from the negative cache rather than rebuilt. A run whose hits are mostly ``hit_unsupported`` is not reusing graphs at all -- it is asking repeatedly for something that will never run fused. Together with ``miss`` these account for every lookup, and ``unsupported`` counts the refusals themselves, so it stays at one per bad configuration however many times that configuration is queried. -.. envvar:: NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES + ``2`` additionally emits a per-lookup ``HIT``/``MISS``/``UNSUPPORTED`` line carrying the full cache key, and a per-execution ``EXEC`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. A level-2 ``UNSUPPORTED`` is a lookup answered from a stored refusal, as opposed to the level-1 event that recorded it; the counter line carries totals, the lookup line carries the key. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. - :Type: ``int`` - :Default: ``500`` - :Description: Ceiling on the number of entries the FusedAttention graph cache keeps per build site, evicting least-recently-used entries to stay under it. There are four build sites (F16 and FP8, forward and backward), each holding a cache of graphs and a cache of configurations cuDNN refused, and the ceiling applies to each of those independently. - - The default is meant to be out of the way of real work rather than tight: a training step reuses a handful of configurations and an inference server with bucketed sequence lengths tens of them, so a few hundred is already more shape diversity than a model exhibits. The ceiling exists for workloads whose key space is effectively unbounded -- a test suite sweeping shapes, or a server keying on something that never repeats -- where an unbounded cache is a slow leak of cuDNN graphs and their execution plans for the life of the process. + By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. - Set to ``0`` to disable the ceiling entirely, for a workload that genuinely has thousands of live configurations and would rather spend the memory than rebuild. Evicting a graph never disturbs one that is executing; execution holds its own reference. + Has negligible overhead when unset. .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 79ba646ca6..1c206264f9 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -259,15 +259,15 @@ "```\n", "NVTE_FUSED_ATTN_CACHE_DEBUG = 0/1/2 # disables/enables graph cache diagnostics\n", "```\n", - "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]`, and there is one per event that happens once per configuration: `BUILD` when a graph is constructed, `PLANS` when its kernels are compiled on first execution, and `UNSUP` when cuDNN declines a configuration. An end-of-run `SUMMARY` gives the totals, per thread and per device, followed by where the build time went:\n", + "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]` and names the build site behind it -- `f16` or `fp8`, then the pass -- and there is one per event that happens once per configuration: `BUILD_GRAPH` when a graph is constructed, `BUILD_PLANS` when its kernels are compiled on first execution, and `UNSUPPORTED` when cuDNN declines a configuration. Each event name is also the counter column it increments, and the two hit columns say which of the cache's maps answered a lookup: `hit_supported` is a cached graph reused, `hit_unsupported` a configuration cuDNN had already refused. An end-of-run `SUMMARY` gives the totals per backend, thread and device, followed by where the build time went:\n", "```\n", - "[FUSED-ATTN-CACHE] pid=1234 | fwd BUILD | tid=0 dev=0 | fwd miss=1, hit=0, built=1, ...\n", - "[FUSED-ATTN-CACHE] pid=1234 | SUMMARY | tid=all dev=all | fwd miss=2, hit=1, built=1, ...\n", - "[FUSED-ATTN-CACHE] pid=1234 | fwd check_support | calls=1 | time= 42.135 ms/call\n", + "[FUSED-ATTN-CACHE] f16 fwd BUILD_GRAPH | tid=0 dev=0 | fwd hit_supported=0, miss=1, build_graph=1, ...\n", + "[FUSED-ATTN-CACHE] f16 SUMMARY | tid=all dev=all | fwd hit_supported=5, miss=1, build_graph=1, ...\n", + "[FUSED-ATTN-CACHE] f16 fwd check_support | calls=1 | time= 42.135 ms/call\n", "```\n", - "The number to read first is `built`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n", + "The number to read first is `build_graph`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n", "\n", - "This flag is supported in both PyTorch and JAX, since the cache it reports on lives in the common C++ layer. By default only rank 0 emits; see [NVTE_FUSED_ATTN_CACHE_DEBUG](../../envvars.rst) for selecting other ranks, and for the companion `NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES` which bounds how many configurations the cache retains." + "This flag is supported in both PyTorch and JAX, since the cache it reports on lives in the common C++ layer. By default only rank 0 emits; see [NVTE_FUSED_ATTN_CACHE_DEBUG](../../envvars.rst) for selecting other ranks." ] }, { diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 3a6cc17d48..b5413d6687 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -275,12 +275,18 @@ def test_dpa_checkpoint(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, None, False, False) -# One [FUSED-ATTN-CACHE] event, as either a counter line ("fwd BUILD") or a level-2 trace -# line ("fwd MISS"). The pass and the event name are all this test reads; the trace line's -# cache key is kept so that distinct keys can be counted. +# One [FUSED-ATTN-CACHE] event, as either a counter line ("f16 fwd BUILD_GRAPH") or a level-2 +# trace line ("f16 fwd MISS"). Both name the build site first, and the backend half of it is +# matched but not counted on: the worker below drives whichever one its dtype selects, and every +# assertion here holds of either. The pass and the event name are what this test reads, plus the +# trace line's cache key, kept so that distinct keys can be counted. Every event name is also the +# counter column it increments, except UNSUPPORTED, which the diagnostics use for both the +# level-1 refusal and the level-2 lookup answered from it -- both count as refusals here. The +# rank prefix is optional because it is emitted only when the launcher exports a rank, which a +# plain subprocess like the worker does not. _CACHE_EVENT = re.compile( - r"\[FUSED-ATTN-CACHE\].*\|\s+(?Pfwd|bwd)\s+" - r"(?PBUILD|PLANS|UNSUP|EXEC|MISS|HIT|NOSUP)\b(?P.*)" + r"\[FUSED-ATTN-CACHE\]\s+(?:rank=\d+\s+\|\s+)?(?Pf16|fp8)\s+(?Pfwd|bwd)\s+" + r"(?PBUILD_GRAPH|BUILD_PLANS|UNSUPPORTED|EXEC|MISS|HIT)\b(?P.*)" ) _CACHE_PHASE = re.compile(r"\[CACHE-TEST\] phase=(?P\w+)") @@ -364,16 +370,16 @@ def count(phase, event, pass_name=pass_name): return events[phase][(pass_name, event)] # The first query builds each pass's graph, and no more than its graph: a support - # query stops at check_support(), leaving the kernel compilation (PLANS) to whoever - # executes it. + # query stops at check_support(), leaving the kernel compilation (BUILD_PLANS) to + # whoever executes it. assert count("query", "MISS") == 1, f"{pass_name}: expected one cold miss{context}" - assert count("query", "BUILD") == 1, f"{pass_name}: expected one build{context}" - assert count("query", "UNSUP") == 0, f"{pass_name}: cuDNN refused the config{context}" - assert count("query", "PLANS") == 0, f"{pass_name}: query compiled kernels{context}" + assert count("query", "BUILD_GRAPH") == 1, f"{pass_name}: expected one build{context}" + assert count("query", "UNSUPPORTED") == 0, f"{pass_name}: cuDNN refused the config{context}" + assert count("query", "BUILD_PLANS") == 0, f"{pass_name}: query compiled kernels{context}" # Asking the identical question again must cost nothing. assert count("requery", "MISS") == 0, f"{pass_name}: repeated query missed{context}" - assert count("requery", "BUILD") == 0, f"{pass_name}: repeated query rebuilt{context}" + assert count("requery", "BUILD_GRAPH") == 0, f"{pass_name}: repeated query rebuilt{context}" assert count("requery", "HIT") >= 1, f"{pass_name}: repeated query never looked{context}" # The execution must find the graph the query left behind -- a miss here is the @@ -382,21 +388,25 @@ def count(phase, event, pass_name=pass_name): assert ( count("exec", "MISS") == 0 ), f"{pass_name}: execution missed the query's graph{context}" - assert count("exec", "BUILD") == 0, f"{pass_name}: execution rebuilt the graph{context}" + assert ( + count("exec", "BUILD_GRAPH") == 0 + ), f"{pass_name}: execution rebuilt the graph{context}" assert count("exec", "EXEC") >= 1, f"{pass_name}: fused attention never ran{context}" - assert count("exec", "PLANS") == 1, f"{pass_name}: expected one plan build{context}" + assert count("exec", "BUILD_PLANS") == 1, f"{pass_name}: expected one plan build{context}" # softmax_scale reaches the graph as a pointer, not as a shape, so the key drops it: # a different scale has to reuse everything, down to the compiled kernels. assert count("rescale", "MISS") == 0, f"{pass_name}: attn_scale changed the key{context}" - assert count("rescale", "BUILD") == 0, f"{pass_name}: attn_scale forced a build{context}" - assert count("rescale", "PLANS") == 0, f"{pass_name}: attn_scale recompiled{context}" + assert ( + count("rescale", "BUILD_GRAPH") == 0 + ), f"{pass_name}: attn_scale forced a build{context}" + assert count("rescale", "BUILD_PLANS") == 0, f"{pass_name}: attn_scale recompiled{context}" assert count("rescale", "EXEC") >= 1, f"{pass_name}: rescaled run did not execute{context}" # max_seqlen is a dimension the graph is built at, so it must miss -- once, for one # new graph, rather than invalidating what is already cached. assert count("reshape", "MISS") == 1, f"{pass_name}: expected one miss{context}" - assert count("reshape", "BUILD") == 1, f"{pass_name}: expected one build{context}" + assert count("reshape", "BUILD_GRAPH") == 1, f"{pass_name}: expected one build{context}" assert ( len(miss_keys["reshape"][pass_name]) == 1 ), f"{pass_name}: more than one new cache key{context}" diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 8a07c14121..451280b2c9 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -55,6 +55,19 @@ void FusedAttnConfig::derive() { (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + // Both layouts describe variable-length sequences inside padded dimensions, so the mask is the + // only thing that tells cuDNN where the real tokens end; without it the graph attends to + // padding. Asserted here so that all four graph builders inherit the rule, and stated as a + // rejection rule in nvte_get_fused_attn_backend_v2 so that a support query answers rather than + // throws. Not conditioned on the cuDNN version: the requirement comes from what the dimensions + // mean, not from what any particular cuDNN can run. + if (is_paged_kv) { + NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); + } + if (is_ragged_q || is_ragged_kv) { + NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal 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); @@ -70,6 +83,23 @@ void FusedAttnConfig::derive() { (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && !is_dropout; + // packed vs dense dimensions for a ragged (THD) graph; SM8x and SM120 require dense, + // BHSD-like dimensions for the Stats/LSE auxiliary tensors and so take the dense path + const int sm_arch = cuda::sm_arch(cuda::current_device()); + uses_packed_ragged_graph = cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; + uses_ragged_stats = is_ragged_q && uses_packed_ragged_graph; + + // sequence lengths the graph is built at + graph_max_seqlen_q = + (is_ragged_q && uses_packed_ragged_graph) ? bucketed_num_tokens_q : max_seqlen_q; + graph_max_seqlen_kv = + (is_ragged_kv && uses_packed_ragged_graph) ? bucketed_num_tokens_kv : max_seqlen_kv; + + // elements per token for each ragged tensor + ragged_offset_mults = RaggedOffsetMultipliers( + layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), + static_cast(head_dim_qk), static_cast(head_dim_v)); + // paged KV dimensions if (is_paged_kv) { if (num_pages_k == 0) { @@ -96,6 +126,10 @@ void FusedAttnConfig::derive() { } FusedAttnConfig FusedAttnConfig::make_cache_key() const { + // Requires a derived config: every normalization below reads a derived field -- is_padding and + // is_causal_bottom_right, the is_ragged_* pair, the graph_max_seqlen_* dimensions, and the + // uses_* flags. A precondition rather than an assert, since all four callers construct their + // GraphInputs first and that constructor asserts it. FusedAttnConfig cache_cfg = *this; // Key the device ID for multi-GPU single-process runs @@ -110,28 +144,25 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { cache_cfg.bottom_right_diagonal = false; } - // 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 (supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_)) { - 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; - // The forward graph keeps the true batch size when it takes the user's cu_seqlens - // directly, since cuDNN reads those [actual_b+1] buffers itself; the backward graph - // converts them and so always buckets. The key has to follow whichever the graph does, - // or it would name a batch size the graph was not built with. See - // derive_f16_bwd_graph_inputs. - const bool bucket_batch = !check_for_forward_support || !cache_cfg.uses_cu_seqlens_directly; - if (bucket_batch) { - cache_cfg.batch_size = cache_cfg.bucketed_batch_size; - } + // Name the sequence lengths the graph is built at rather than the ones the caller asked about, + // so that every shape falling in the same bucket lands on the same entry. The two are equal + // unless a ragged layout is packed, which is why this is unconditional. Stated after the + // bottom_right_diagonal rule above, which is about the real geometry of the attention mask and + // would read bucketed token counts as sequence lengths if it ran after the substitution. + cache_cfg.max_seqlen_q = cache_cfg.graph_max_seqlen_q; + cache_cfg.max_seqlen_kv = cache_cfg.graph_max_seqlen_kv; + + // Bucket the THD (ragged) batch, and drop the token counts the bucketing has replaced + if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_packed_ragged_graph) { + cache_cfg.num_tokens_q = 0; + cache_cfg.num_tokens_kv = 0; + // The forward graph keeps the true batch size when it takes the user's cu_seqlens + // directly, since cuDNN reads those [actual_b+1] buffers itself; the backward graph + // converts them and so always buckets. The key has to follow whichever the graph does, + // or it would name a batch size the graph was not built with. See F16BwdGraphInputs. + const bool bucket_batch = !check_for_forward_support || !cache_cfg.uses_cu_seqlens_directly; + if (bucket_batch) { + cache_cfg.batch_size = cache_cfg.bucketed_batch_size; } } @@ -146,14 +177,15 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { // Restrict each direction's key to the fields its graph actually consumes, so // no redundant graphs are built and no cache misses either - if (check_for_forward_support) { + if (check_for_forward_support && !check_for_backward_support) { 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 { + } + if (check_for_backward_support && !check_for_forward_support) { 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 d116893adf..44e0ef457b 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -15,19 +15,11 @@ #include "common/common.h" #include "transformer_engine/fused_attn.h" +#include "utils.h" namespace transformer_engine { namespace fused_attn { -// Whether a ragged (THD) graph can be built at packed token-count dimensions with ragged -// Stats/LSE. SM8x and SM120 require dense, BHSD-like dimensions at max_seqlen for the auxiliary -// tensors instead. Graph construction, auxiliary-tensor allocation and make_cache_key() all -// answer this question, and a disagreement between them would key a graph by dimensions it was -// not built with, so they share this one definition. -inline constexpr bool supports_packed_ragged_graph(size_t cudnn_runtime_version, int sm_arch) { - return cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; -} - struct FusedAttnConfig { // basic attention settings bool is_training = true; @@ -98,12 +90,40 @@ struct FusedAttnConfig { // so this exists to let those consumers assert rather than trust. Not a cached-result marker: // derive() recomputes unconditionally, so a config whose inputs change can simply be re-derived. bool is_derived = false; - // THD batch/token counts; make_cache_key() folds these into batch_size/max_seqlen_*. + // THD batch/token counts, the raw buckets. The graph dimensions built out of them are + // graph_max_seqlen_* below and, because the batch is direction-dependent, F16FwdGraphInputs::b. 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; + // Whether a ragged (THD) graph is built at packed token-count dimensions with ragged Stats/LSE, + // rather than at dense max_seqlen ones. Held here rather than asked for at each of the places + // that need it -- graph_max_seqlen_* below, make_cache_key()'s batch, and the two GraphInputs -- + // because the key and the graph have to be built at the same dimensions, and two independent + // queries are two chances to disagree. Unlike the flags above, this one depends on the device as + // well as the cuDNN version, so a config carries the answer for the device it was derived on; + // every entry point derives immediately before use, and the cache key records device_id. + bool uses_packed_ragged_graph = false; + // Whether the graph's Stats/LSE tensor is the packed, token-indexed one. Ragged Q is necessary + // but not sufficient, since the packed representation also needs an architecture that supports + // it. Derived because three unrelated places read it -- the graph build, the pointer binding at + // execution, and the Stats/Max shapes reported back to the framework -- and they are describing + // one buffer, so they cannot be allowed to disagree about its shape. + bool uses_ragged_stats = false; + // The sequence lengths the graph is built at: max_seqlen_* for a dense graph, and the bucketed + // token counts where a ragged layout is packed. Held here because the cache key has to name the + // dimensions the graph was built with -- a key that says otherwise is a hit on a graph of the + // wrong shape -- and stating the substitution once is what keeps make_cache_key() and the graph + // builders from drifting. Both passes build at the same sequence lengths; the batch size is the + // one dimension they disagree on, so it stays with the direction that knows, in + // F16FwdGraphInputs and F16BwdGraphInputs. + size_t graph_max_seqlen_q = 0; + size_t graph_max_seqlen_kv = 0; + // Elements per token for each ragged tensor, from the layout group and the head dimensions. + // Shared with the cu_seqlens_padded_to_offsets kernel, so the offsets the graph is told to + // expect and the offsets that are written cannot drift apart. + RaggedOffsetMultipliers ragged_offset_mults; // 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; @@ -208,9 +228,15 @@ struct FusedAttnConfig { // configuration is supported does not modify it. Nothing further in is expected to derive // again, and check_derived() is what holds them to that. Idempotent, so a config that is // derived and then re-derived is unharmed. + // + // Throws for combinations of input fields that no graph can serve, so that all four graph + // builders inherit the rule from one place. Those same combinations are stated as rejection + // rules in nvte_get_fused_attn_backend_v2(), ahead of its derive() call, so that asking whether + // such a configuration is supported gets an answer instead of an exception. void derive(); // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. + // Requires a config that has been through derive(), whose fields the normalizations read. // 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. @@ -222,7 +248,9 @@ struct FusedAttnConfig { // q_format reads as zero, which is a legal value that yields a graph of the wrong shape and a // key that collides with unrelated configs. Deriving happens at the library's entry points rather // than here, where it would be needed, so this is what keeps a new path into the builders from -// quietly skipping it. +// quietly skipping it. It catches a config that was never derived and nothing else: a config +// derived and then edited passes, so callers that change an input field re-derive rather than rely +// on this, which derive() being idempotent makes cheap. inline void check_derived(const FusedAttnConfig &cfg) { NVTE_CHECK(cfg.is_derived, "FusedAttnConfig reached a graph build with its derived fields unset. Every config " diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index c47dc9b603..a7c5e25ef9 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -240,13 +240,14 @@ void set_message(const char **message, std::string reason) { *message = fused_attn_backend_message_buffer.c_str(); } -// Records `reason` if `rejected`, and reports it, so that an early rejection reads as the one -// statement it is: `if (set_message_if(cond, message, "why")) return NVTE_No_Backend;`. The -// reason is built whether or not it is used, which is why it stays a plain string here -- these -// are short literals on a path that goes on to build cuDNN graphs. -bool set_message_if(bool rejected, const char **message, std::string reason) { - if (rejected) set_message(message, std::move(reason)); - return rejected; +// Records `reason` and answers with the backend that means "none", so that a rejection reads as +// the one statement it is: `if (cond) return reject(message, "why");`. Every rejection in +// nvte_get_fused_attn_backend_v2 goes through here, which is what keeps a reason attached to +// each: the value cannot be produced without one. nodiscard because dropping the value would +// leave the message set and the rejection unreturned, and the function would carry on. +[[nodiscard]] NVTE_Fused_Attn_Backend reject(const char **message, std::string reason) { + set_message(message, std::move(reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } // namespace @@ -267,13 +268,18 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi // inputs, so make_cache_key() lands on the same entry. Deriving is idempotent, so re-deriving // an already-derived config here changes nothing. FusedAttnConfig cfg = *get_fused_attn_config(config); - cfg.derive(); set_message(message, ""); cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); const auto qkv_format = nvte_get_qkv_format(cfg.qkv_layout); const auto layout_group = nvte_get_qkv_layout_group(cfg.qkv_layout); const auto cudnn_runtime_version = cudnnGetVersion(); + // Read from attn_mask_type rather than from cfg.is_padding, because the two rules that need it + // are stated before derive() runs; see the derive() call below. + const bool has_padding_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; // THD + 64-bit ragged offsets require cuDNN >= 9.5 const bool requires_64bit_ragged_offset = @@ -281,26 +287,33 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi 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 (set_message_if(requires_64bit_ragged_offset && cudnn_runtime_version < 90500, message, - "Configuration requires 64-bit ragged offsets, which require " - "cuDNN >= 9.5.")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { + return reject(message, + "Configuration requires 64-bit ragged offsets, which require cuDNN >= 9.5."); } // THD requires padding-style mask - if (set_message_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, - message, - "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask.")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (qkv_format == NVTE_QKV_Format::NVTE_THD && !has_padding_mask) { + return reject( + message, + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + } + + // Paged KV requires padding-style mask, for the same reason THD does: the graph is built at + // padded dimensions and the mask is what tells cuDNN where the real tokens end. + if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && !has_padding_mask) { + return reject(message, + "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); } + // Derived here rather than above, so that the two rules stated above are answered rather than + // thrown. Both are invariants derive() asserts, and an assertion that fired first would leave + // this function no chance to report them as an unsupported configuration. + cfg.derive(); + // Ragged Q/KV requires sm90+, the rule the hand-written support matrix this function replaced // carried as `qkv_format == NVTE_THD && sm_arch_ >= 90`. Below sm90 the only graph we can build - // is the dense max_seqlen one -- supports_packed_ragged_graph() is false -- so SDPA_backward + // is the dense max_seqlen one -- cfg.uses_packed_ragged_graph is false -- so SDPA_backward // never gets max_total_seq_len_q/kv and its dQ/dK/dV come back wrong. // // This is ours to state because it is a wrong-result rejection, and check_support answers a @@ -316,16 +329,13 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi // // sm120 takes that same dense path and is left enabled, as it was before this refactor; // whether it has the same problem is a separate question from restoring the sm90 rule. - if (set_message_if((cfg.is_ragged_q || cfg.is_ragged_kv) && - cuda::sm_arch(cuda::current_device()) < 90, - message, "Ragged (THD) Q or KV requires compute capability 9.0 or higher.")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if ((cfg.is_ragged_q || cfg.is_ragged_kv) && cuda::sm_arch(cuda::current_device()) < 90) { + return reject(message, "Ragged (THD) Q or KV requires compute capability 9.0 or higher."); } // TE's cuDNN fused-attention graph does not represent pre-scale bias. - if (set_message_if(cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS, message, - "Fused attention does not support pre-scale bias.")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) { + return reject(message, "Fused attention does not support pre-scale bias."); } const bool is_fp8 = @@ -334,31 +344,21 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); if (is_fp8) { - if (set_message_if(cfg.return_max_logit, message, - "FP8 fused attention does not support return_max_logit=True.")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (cfg.return_max_logit) { + return reject(message, "FP8 fused attention does not support return_max_logit=True."); } - if (set_message_if(qkv_format != NVTE_QKV_Format::NVTE_BSHD && - qkv_format != NVTE_QKV_Format::NVTE_SBHD && - qkv_format != NVTE_QKV_Format::NVTE_BHSD, - message, - "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + - std::to_string(static_cast(qkv_format)) + ".")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + if (qkv_format != NVTE_QKV_Format::NVTE_BSHD && qkv_format != NVTE_QKV_Format::NVTE_SBHD && + qkv_format != NVTE_QKV_Format::NVTE_BHSD) { + return reject(message, "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(qkv_format)) + "."); } if (cfg.check_for_forward_support) { 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 (!fwd_reason.empty()) return reject(message, std::move(fwd_reason)); } if (cfg.is_training && cfg.check_for_backward_support) { std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); - if (!bwd_reason.empty()) { - set_message(message, std::move(bwd_reason)); - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; - } + if (!bwd_reason.empty()) return reject(message, std::move(bwd_reason)); } return NVTE_Fused_Attn_Backend::NVTE_FP8; } @@ -370,36 +370,26 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi // check_support is the authority now, so the guard is gone; it needs to come back as an // explicit rejection here, like the CUDA-graph one below, if that bug is a wrong-result // bug rather than a support gap check_support reports for itself. - if (set_message_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, - message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN.")) { - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + 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) { + return reject(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); } if (cfg.check_for_forward_support) { 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 (!fwd_reason.empty()) return reject(message, std::move(fwd_reason)); } if (cfg.is_training && cfg.check_for_backward_support) { 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; - } + if (!bwd_reason.empty()) return reject(message, std::move(bwd_reason)); } return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg.qkv_dtype) + " ."); - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + return reject(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg.qkv_dtype) + " ."); } // select a backend for fused attention @@ -448,7 +438,36 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( /*message=*/nullptr); } -// fused attention forward +// Fused attention forward: derive the config, ask the selector which backend can run it, and run +// that backend's implementation. +// +// Support is decided by building the graph rather than by consulting a table of rules, and the +// support query and the execution path reach the same cache through the same accessor. That is +// what the HIT below means: by the time a backend has been selected, the entry the implementation +// needs has already been built and inserted by the probe that selected it, so what was checked is +// what runs. The rules the selector does state for itself are the ones cuDNN cannot answer, +// because they are about whether the graph computes what was asked for rather than whether cuDNN +// can run it. +// +// nvte_fused_attn_fwd_v2 +// | +// +-- cfg = p.make_config(), which sets check_for_forward_support; cfg.derive() +// | +// +-- nvte_get_fused_attn_backend_v2 the support query +// | | +// | +-- TE's own rules: THD and paged KV need a padding mask, no pre-scale bias, +// | | ragged Q/KV needs sm90+, the cuDNN 9.15-and-older CUDA-graph bug +// | | `-- reject -> NVTE_No_Backend + reason -> the NVTE_ERROR below +// | | +// | `-- is_supported_f16_fwd / is_supported_fp8_fwd +// | `-- f16_fwd_cached_graph(): builds and inserts the entry, or throws +// | UnsupportedGraph, whose message becomes the reason for the refusal +// | +// `-- fused_attn_arbitrary_seqlen_fwd -> ..._fwd_impl the selected backend +// | +// +-- f16_fwd_cached_graph() HIT: the entry the query above just built +// +-- ensure_plans_built() the kernel compilation, once per entry +// `-- bind device pointers, execute() void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; @@ -551,7 +570,9 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso nvte_fused_attn_fwd_v2(reinterpret_cast(&p)); } -// fused attention backward +// Fused attention backward. The same shape as nvte_fused_attn_fwd_v2, which sketches the path from +// an entry point through the selector to the cache; the only differences are that the config asks +// the selector for backward support and that the backward builders are the ones the probe runs. void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { NVTE_API_CALL(nvte_fused_attn_bwd_v2); using namespace transformer_engine; 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 da3bf7a875..42130de8ff 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 @@ -25,6 +25,11 @@ namespace fused_attn { namespace fe = cudnn_frontend; +// Every graph-cache event raised here names the build site it came from. This file is the f16 +// arbitrary-seqlen backend throughout; only the pass differs between call sites. +using graph_cache_debug::Backend; +using graph_cache_debug::Pass; + using SdpaF16FwdGraphAndTensors = std::tuple, std::shared_ptr, // Q @@ -48,102 +53,54 @@ using SdpaF16FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// What the forward graph is built from beyond the config's own fields: the dimensions ragged -// layouts bucket, and the choices that depend on the cuDNN runtime version or the SM -// architecture. The build and the execution have to reach the same answer for every one of -// these -- otherwise the graph is built for different dimensions than the pointers bound to it -// describe, or with a ragged offset width the offsets are not written in -- so they are derived -// once, by derive_f16_fwd_graph_inputs, and handed to both. +// What the forward graph is built from that the config cannot say on its own, because the answer +// depends on the direction: the batch size, and the width the ragged offsets are written in. The +// sequence lengths are not here -- both passes build at cfg.graph_max_seqlen_* -- and neither is +// anything else a backward graph would answer the same way. The build and the execution have to +// agree on all of it, otherwise the graph is built for different dimensions than the pointers +// bound to it describe, or with a ragged offset width the offsets are not written in, so it is +// derived once, here, and handed to both. struct F16FwdGraphInputs { - // Dimensions the graph is built at. Ragged layouts substitute bucketed token counts for - // max_seqlen (and, unless cu_seqlens are passed to cuDNN directly, a bucketed batch size) - // so that one graph serves every shape that falls in the same bucket. - int64_t b; - int64_t s_q; - int64_t s_kv; + // Everything below is arithmetic on an already-derived cfg. Configurations no graph can serve + // are rejected before this point: FusedAttnConfig::derive() asserts them, and + // nvte_get_fused_attn_backend_v2 states them as rules so a support query can answer for them. + explicit F16FwdGraphInputs(const FusedAttnConfig &cfg); + + // The batch size the graph is built at: bucketed for a packed ragged layout, so that one graph + // serves every batch in the same bucket, except when cu_seqlens go to cuDNN directly. + int64_t b = 0; // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever // the bucketing above did to `b`. - int64_t actual_b; - // Whether this architecture takes the packed token-count representation above; see - // supports_packed_ragged_graph. Carried here so the graph build reads the same answer the - // dimensions were derived from rather than querying the device again. - bool use_packed_ragged_graph; - bool use_ragged_stats; - DType ragged_offset_type; - RaggedOffsetMultipliers offset_mults; + int64_t actual_b = 0; + DType ragged_offset_type = DType::kInt32; }; -// Derives the above, and rejects configurations that no graph can serve. Those rejections -// depend on combinations of fields rather than any single one, so they cannot live in the -// config's own validation; running them here is what lets a support query answer for them -// without building anything. -static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) { +F16FwdGraphInputs::F16FwdGraphInputs(const FusedAttnConfig &cfg) { check_derived(cfg); - const bool is_padding = cfg.is_padding; const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; const auto cudnn_runtime_version = cudnnGetVersion(); - const int sm_arch_ = cuda::sm_arch(cuda::current_device()); - - if (cfg.is_paged_kv) { - NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); - } - - const bool use_packed_ragged_graph = - supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); - int64_t b = static_cast(cfg.batch_size); - int64_t s_q = static_cast(cfg.max_seqlen_q); - int64_t s_kv = static_cast(cfg.max_seqlen_kv); + b = static_cast(cfg.batch_size); // keep original batch size because cu_seqlens are created with [b+1] shape - const int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // SM8x and SM120 need dense, BHSD-like dimensions/strides at max_seqlen: on SM120 the cuDNN - // support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] as interleaved and - // rejects them. The ragged offsets still provide the variable-length boundaries either way. - if (use_packed_ragged_graph) { - // 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. - // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: - // cuDNN reads the user's [actual_b+1] cu_seqlens buffers, so a quantized batch - // would read out of bounds. - if (!use_cu_seqlens_directly) { - b = static_cast(cfg.bucketed_batch_size); - } - s_q = is_ragged_q ? static_cast(cfg.bucketed_num_tokens_q) : s_q; - s_kv = is_ragged_kv ? static_cast(cfg.bucketed_num_tokens_kv) : s_kv; - } + actual_b = b; + // Replace the batch size with the bucketed one so the graph is static within its bucket, the + // same reason cfg.graph_max_seqlen_* replaces the sequence lengths. When passing cu_seqlens* + // directly to cuDNN SDPA, keep the true batch size: cuDNN reads the user's [actual_b+1] + // cu_seqlens buffers, so a quantized batch would read out of bounds. + if ((is_ragged_q || is_ragged_kv) && cfg.uses_packed_ragged_graph && !use_cu_seqlens_directly) { + b = static_cast(cfg.bucketed_batch_size); } - const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph; - const DType ragged_offset_type = + ragged_offset_type = use_cu_seqlens_directly ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. : (cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32); - // Ragged offset multipliers (elements per token); shared with the legacy conversion - // kernel (cu_seqlens_padded_to_offsets) so the two paths cannot drift apart. - const RaggedOffsetMultipliers offset_mults( - nvte_get_qkv_layout_group(cfg.qkv_layout), static_cast(cfg.num_attn_heads), - static_cast(cfg.num_gqa_groups), static_cast(cfg.head_dim_qk), - static_cast(cfg.head_dim_v)); - - // Field order must match F16FwdGraphInputs; one per line so that it can be checked by eye. - return F16FwdGraphInputs{ - b, - s_q, - s_kv, - actual_b, - use_packed_ragged_graph, - use_ragged_stats, - ragged_offset_type, - offset_mults, - }; } // Constructs the forward graph for one cache key, and only constructs it: whether cuDNN will run -// it is settled by the caller, in get_or_build_cached_graph(), which is also where the plan build +// it is settled by the caller, in build_or_get_cached_graph(), which is also where the plan build // eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every call // that does need one now sits on the other side of that boundary. // @@ -152,8 +109,8 @@ static F16FwdGraphInputs derive_f16_fwd_graph_inputs(const FusedAttnConfig &cfg) static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig &cfg, const F16FwdGraphInputs &in) { const int64_t b = in.b; - const int64_t s_q = in.s_q; - const int64_t s_kv = in.s_kv; + const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); const cudnn_frontend::DataType_t tensorType = get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); const int64_t h = static_cast(cfg.num_attn_heads); @@ -193,9 +150,9 @@ static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig const bool is_ragged_kv = cfg.is_ragged_kv; const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; const auto cudnn_runtime_version = cudnnGetVersion(); - const bool use_ragged_stats = in.use_ragged_stats; + const bool use_ragged_stats = cfg.uses_ragged_stats; const DType ragged_offset_type = in.ragged_offset_type; - const RaggedOffsetMultipliers offset_mults = in.offset_mults; + const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; const bool generate_stats = true; // Always return stats auto mha_graph = std::make_shared(); @@ -466,7 +423,7 @@ static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig static std::shared_ptr> f16_fwd_cached_graph( const FusedAttnConfig &cfg, const F16FwdGraphInputs &in, cudnnHandle_t handle) { static GraphCache cache; - return get_or_build_cached_graph(cache, cfg.make_cache_key(), "fwd", handle, + return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::F16, Pass::Fwd, handle, [&] { return build_sdpa_f16_fwd_graph(cfg, in); }); } @@ -480,14 +437,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( using namespace transformer_engine; // Derived once and handed to the cache, which passes them to the graph build, so that the - // graph and the pointers bound to it below cannot be decided differently. Also where an - // unserviceable configuration is rejected. - const F16FwdGraphInputs in = derive_f16_fwd_graph_inputs(cfg); + // graph and the pointers bound to it below cannot be decided differently. + const F16FwdGraphInputs in(cfg); const int64_t b = in.b; const int64_t actual_b = in.actual_b; - const bool use_ragged_stats = in.use_ragged_stats; const DType ragged_offset_type = in.ragged_offset_type; - const RaggedOffsetMultipliers offset_mults = in.offset_mults; + const bool use_ragged_stats = cfg.uses_ragged_stats; + const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; const bool return_max_logit = cfg.return_max_logit; // Not const: bound into the variant pack by address as a pass-by-value graph input. @@ -511,7 +467,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( dropout_seed, dropout_offset] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built("fwd", *cache_entry); + ensure_plans_built(Backend::F16, Pass::Fwd, *cache_entry); // 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. @@ -540,7 +496,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"); + graph_cache_debug::record_exec(Backend::F16, Pass::Fwd); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -681,63 +637,35 @@ using SdpaF16BwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The backward equivalent of F16FwdGraphInputs; see there for why these are derived once and -// shared. The backward graph reads no page table and passes no cu_seqlens straight through, so -// it needs neither the paged-attention check nor the ragged offset multipliers. +// The backward equivalent of F16FwdGraphInputs; see there for why these two dimensions are the +// only ones that cannot live on the config, and for why they are derived once and shared. struct F16BwdGraphInputs { - int64_t b; - int64_t s_q; - int64_t s_kv; - int64_t actual_b; - bool use_packed_ragged_graph; - bool use_ragged_stats; - DType ragged_offset_type; + explicit F16BwdGraphInputs(const FusedAttnConfig &cfg); + + int64_t b = 0; + int64_t actual_b = 0; + DType ragged_offset_type = DType::kInt32; }; -// The backward counterpart of derive_f16_fwd_graph_inputs; see there for what the rejections are -// doing here and why a support query can answer for them without building a graph. -static F16BwdGraphInputs derive_f16_bwd_graph_inputs(const FusedAttnConfig &cfg) { +F16BwdGraphInputs::F16BwdGraphInputs(const FusedAttnConfig &cfg) { check_derived(cfg); - const bool is_padding = cfg.is_padding; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); - const int sm_arch_ = cuda::sm_arch(cuda::current_device()); - const bool use_packed_ragged_graph = - supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); - - int64_t b = static_cast(cfg.batch_size); - int64_t s_q = static_cast(cfg.max_seqlen_q); - int64_t s_kv = static_cast(cfg.max_seqlen_kv); + b = static_cast(cfg.batch_size); // keep original batch size because cu_seqlens are created with [b+1] shape - const int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // SM8x and SM120 require dense, BHSD-like strides at max_seqlen (see fwd). - if (use_packed_ragged_graph) { - // 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. - // The batch is bucketed unconditionally here, where the forward pass guards it: only - // the forward graph can be handed the user's cu_seqlens buffers directly, and it is - // their [actual_b+1] length that a quantized batch would overrun. The backward graph - // always reads converted seqlens out of our own workspace, so nothing here is sized by - // the true batch. make_cache_key() splits on the pass for this reason as well. - b = static_cast(cfg.bucketed_batch_size); - s_q = is_ragged_q ? static_cast(cfg.bucketed_num_tokens_q) : s_q; - s_kv = is_ragged_kv ? static_cast(cfg.bucketed_num_tokens_kv) : s_kv; - } + actual_b = b; + // The batch is bucketed unconditionally here, where the forward pass guards it: only the + // forward graph can be handed the user's cu_seqlens buffers directly, and it is their + // [actual_b+1] length that a quantized batch would overrun. The backward graph always reads + // converted seqlens out of our own workspace, so nothing here is sized by the true batch. + // make_cache_key() splits on the pass for this reason as well. + if ((cfg.is_ragged_q || cfg.is_ragged_kv) && cfg.uses_packed_ragged_graph) { + b = static_cast(cfg.bucketed_batch_size); } - const bool use_ragged_stats = is_ragged_q && use_packed_ragged_graph; // 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; - - // Field order must match F16BwdGraphInputs; one per line so that it can be checked by eye. - return F16BwdGraphInputs{ - b, s_q, s_kv, actual_b, use_packed_ragged_graph, use_ragged_stats, ragged_offset_type, - }; + ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; } // The backward counterpart of build_sdpa_f16_fwd_graph; see there for why it constructs the graph @@ -748,8 +676,8 @@ static F16BwdGraphInputs derive_f16_bwd_graph_inputs(const FusedAttnConfig &cfg) static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig &cfg, const F16BwdGraphInputs &in) { const int64_t b = in.b; - const int64_t s_q = in.s_q; - const int64_t s_kv = in.s_kv; + const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); const cudnn_frontend::DataType_t tensorType = get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); const int64_t h = static_cast(cfg.num_attn_heads); @@ -780,8 +708,8 @@ static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); - const bool use_packed_ragged_graph = in.use_packed_ragged_graph; - const bool use_ragged_stats = in.use_ragged_stats; + const bool use_packed_ragged_graph = cfg.uses_packed_ragged_graph; + const bool use_ragged_stats = cfg.uses_ragged_stats; const DType ragged_offset_type = in.ragged_offset_type; auto mha_graph = std::make_shared(); @@ -1016,7 +944,7 @@ static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig static std::shared_ptr> f16_bwd_cached_graph( const FusedAttnConfig &cfg, const F16BwdGraphInputs &in, cudnnHandle_t handle) { static GraphCache cache; - return get_or_build_cached_graph(cache, cfg.make_cache_key(), "bwd", handle, + return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::F16, Pass::Bwd, handle, [&] { return build_sdpa_f16_bwd_graph(cfg, in); }); } @@ -1031,21 +959,15 @@ void fused_attn_arbitrary_seqlen_bwd_impl( using namespace transformer_engine; // Derived once and handed to the cache, which passes them to the graph build, so that the - // graph and the pointers bound to it below cannot be decided differently. Also where an - // unserviceable configuration is rejected. - const F16BwdGraphInputs in = derive_f16_bwd_graph_inputs(cfg); + // graph and the pointers bound to it below cannot be decided differently. + const F16BwdGraphInputs in(cfg); const int64_t b = in.b; const int64_t actual_b = in.actual_b; - const bool use_ragged_stats = in.use_ragged_stats; const DType ragged_offset_type = in.ragged_offset_type; + const bool use_ragged_stats = cfg.uses_ragged_stats; - const int64_t h = static_cast(cfg.num_attn_heads); - const int64_t hg = static_cast(cfg.num_gqa_groups); - const int64_t d_qk = static_cast(cfg.head_dim_qk); - const int64_t d_v = static_cast(cfg.head_dim_v); // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); const bool is_padding = cfg.is_padding; const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); @@ -1060,7 +982,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( dropout_seed, dropout_offset] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built("bwd", *cache_entry); + ensure_plans_built(Backend::F16, Pass::Bwd, *cache_entry); // 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. @@ -1084,7 +1006,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"); + graph_cache_debug::record_exec(Backend::F16, Pass::Bwd); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1149,10 +1071,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const RaggedOffsetMultipliers offset_mults(nvte_get_qkv_layout_group(qkv_layout), h, hg, d_qk, - d_v); cu_seqlens_padded_to_offsets<<>>( - offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), + cfg.ragged_offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1222,9 +1142,6 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; @@ -1234,11 +1151,9 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i size_t i = 0; if (Aux_CTX_Tensors->size == 0) { - const auto cudnn_runtime_version = cudnnGetVersion(); - // These have to match the shape the forward graph declares for Stats and Max, which is - // packed only where the architecture supports it; see derive_f16_fwd_graph_inputs. - const bool use_ragged_stats = - cfg.is_ragged_q && supports_packed_ragged_graph(cudnn_runtime_version, sm_arch_); + // These have to match the shape the forward graph declares for Stats and Max, which is why + // both read the same derived field rather than recomputing the condition. + const bool use_ragged_stats = cfg.uses_ragged_stats; Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; @@ -1421,9 +1336,10 @@ 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.check_for_forward_support = true; + graph_cfg.check_for_backward_support = false; try { - const fused_attn::F16FwdGraphInputs in = fused_attn::derive_f16_fwd_graph_inputs(graph_cfg); + const fused_attn::F16FwdGraphInputs in(graph_cfg); fused_attn::f16_fwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception &e) { @@ -1437,9 +1353,10 @@ 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.check_for_forward_support = false; + graph_cfg.check_for_backward_support = true; try { - const fused_attn::F16BwdGraphInputs in = fused_attn::derive_f16_bwd_graph_inputs(graph_cfg); + const fused_attn::F16BwdGraphInputs in(graph_cfg); fused_attn::f16_bwd_cached_graph(graph_cfg, in, handle); return ""; } 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 4214be8614..a4211275a3 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -20,6 +20,11 @@ namespace fused_attn { using namespace transformer_engine; namespace fe = cudnn_frontend; +// Every graph-cache event raised here names the build site it came from. This file is the fp8 +// backend throughout; only the pass differs between call sites. +using graph_cache_debug::Backend; +using graph_cache_debug::Pass; + // fused attention FWD FP8 with FE 1.0+ using SdpaFp8FwdGraphAndTensors = std::tuple, @@ -49,16 +54,17 @@ using SdpaFp8FwdGraphAndTensors = // has, and so which pointers the variant pack has to bind -- the build and the execution cannot // answer them differently, which is why they are derived once, here, for both. struct Fp8FwdGraphInputs { - bool is_delayed_scaling; - bool is_current_scaling; - bool is_mxfp8; - bool use_cu_seqlens_directly; + // Also where what FP8 cannot serve is rejected. Unlike the F16 path there is no bucketing to + // do, because FP8 has no ragged/THD support: the graph's shapes are exactly the config's. + explicit Fp8FwdGraphInputs(const FusedAttnConfig& cfg); + + bool is_delayed_scaling = false; + bool is_current_scaling = false; + bool is_mxfp8 = false; + bool use_cu_seqlens_directly = false; }; -// Derives the above and rejects what FP8 cannot serve. Unlike the F16 path there is no -// bucketing to do, because FP8 has no ragged/THD support: the graph's shapes are exactly the -// config's. -static Fp8FwdGraphInputs derive_fp8_fwd_graph_inputs(const FusedAttnConfig& cfg) { +Fp8FwdGraphInputs::Fp8FwdGraphInputs(const FusedAttnConfig& cfg) { check_derived(cfg); const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t o_tensor_type = @@ -70,15 +76,15 @@ static Fp8FwdGraphInputs derive_fp8_fwd_graph_inputs(const FusedAttnConfig& cfg) 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!"); - const bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - const bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - const bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || + o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); + is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); NVTE_CHECK( is_delayed_scaling || is_current_scaling || is_mxfp8, "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); @@ -90,7 +96,7 @@ static Fp8FwdGraphInputs derive_fp8_fwd_graph_inputs(const FusedAttnConfig& cfg) // the F16 path, the FP8 path has no THD/ragged-offset support, so only the // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the // needed versions of cuDNN backend and frontend are higher than for F16.) - const bool use_cu_seqlens_directly = + use_cu_seqlens_directly = // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). // Note: For the Python API, 1.27 is required. CUDNN_FRONTEND_VERSION >= 12600 && @@ -102,18 +108,10 @@ static Fp8FwdGraphInputs derive_fp8_fwd_graph_inputs(const FusedAttnConfig& cfg) // 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; - - // Field order must match Fp8FwdGraphInputs; one per line so that it can be checked by eye. - return Fp8FwdGraphInputs{ - is_delayed_scaling, - is_current_scaling, - is_mxfp8, - use_cu_seqlens_directly, - }; } // Constructs the forward FP8 graph for one cache key, and only constructs it: whether cuDNN will -// run it is settled by the caller, in get_or_build_cached_graph(), which is also where the plan +// run it is settled by the caller, in build_or_get_cached_graph(), which is also where the plan // build eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every // call that does need one now sits on the other side of that boundary. // @@ -129,8 +127,8 @@ static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& 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 s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_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 int64_t window_size_left = cfg.window_size_left; @@ -403,7 +401,7 @@ static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& static std::shared_ptr> fp8_fwd_cached_graph( const FusedAttnConfig& cfg, const Fp8FwdGraphInputs& in, cudnnHandle_t handle) { static GraphCache cache; - return get_or_build_cached_graph(cache, cfg.make_cache_key(), "fwd", handle, + return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::FP8, Pass::Fwd, handle, [&] { return build_sdpa_fp8_fwd_graph(cfg, in); }); } @@ -420,7 +418,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de // Derived once and handed to the cache, which passes them to the graph build, so that the // graph and the pointers bound to it below cannot be decided differently. Also where an // unserviceable configuration is rejected. - const Fp8FwdGraphInputs in = derive_fp8_fwd_graph_inputs(cfg); + const Fp8FwdGraphInputs in(cfg); const bool is_delayed_scaling = in.is_delayed_scaling; const bool is_current_scaling = in.is_current_scaling; const bool use_cu_seqlens_directly = in.use_cu_seqlens_directly; @@ -440,7 +438,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de dropout_offset] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built("fwd", *cache_entry); + ensure_plans_built(Backend::FP8, Pass::Fwd, *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -452,7 +450,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"); + graph_cache_debug::record_exec(Backend::FP8, Pass::Fwd); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -574,15 +572,17 @@ using SdpaFp8BwdGraphAndTensors = // rather than O's, since backward is what writes those. is_O_in_F16 additionally selects whether // O has to be descaled on the way in. struct Fp8BwdGraphInputs { - bool is_delayed_scaling; - bool is_current_scaling; - bool is_mxfp8; - bool is_O_in_F16; + // The backward counterpart of Fp8FwdGraphInputs' constructor; see there for the rejections and + // for why the graph's shapes are simply the config's. + explicit Fp8BwdGraphInputs(const FusedAttnConfig& cfg); + + bool is_delayed_scaling = false; + bool is_current_scaling = false; + bool is_mxfp8 = false; + bool is_O_in_F16 = false; }; -// The backward counterpart of derive_fp8_fwd_graph_inputs; see there for the rejections and for -// why the graph's shapes are simply the config's. -static Fp8BwdGraphInputs derive_fp8_bwd_graph_inputs(const FusedAttnConfig& cfg) { +Fp8BwdGraphInputs::Fp8BwdGraphInputs(const FusedAttnConfig& cfg) { check_derived(cfg); const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t o_tensor_type = @@ -595,31 +595,23 @@ static Fp8BwdGraphInputs derive_fp8_bwd_graph_inputs(const FusedAttnConfig& cfg) 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!"); - const bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - const bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - const bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || + dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); + is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || + dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || + dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); NVTE_CHECK( is_delayed_scaling || is_current_scaling || is_mxfp8, "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - const bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - - // Field order must match Fp8BwdGraphInputs; one per line so that it can be checked by eye. - return Fp8BwdGraphInputs{ - is_delayed_scaling, - is_current_scaling, - is_mxfp8, - is_O_in_F16, - }; + is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); } // The backward counterpart of build_sdpa_fp8_fwd_graph; see there for why it constructs the graph @@ -638,8 +630,8 @@ static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& 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 s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_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 int64_t window_size_left = cfg.window_size_left; @@ -1045,7 +1037,7 @@ static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& static std::shared_ptr> fp8_bwd_cached_graph( const FusedAttnConfig& cfg, const Fp8BwdGraphInputs& in, cudnnHandle_t handle) { static GraphCache cache; - return get_or_build_cached_graph(cache, cfg.make_cache_key(), "bwd", handle, + return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::FP8, Pass::Bwd, handle, [&] { return build_sdpa_fp8_bwd_graph(cfg, in); }); } @@ -1066,7 +1058,7 @@ void fused_attn_fp8_bwd_impl( // Derived once and handed to the cache, which passes them to the graph build, so that the // graph and the pointers bound to it below cannot be decided differently. Also where an // unserviceable configuration is rejected. - const Fp8BwdGraphInputs in = derive_fp8_bwd_graph_inputs(cfg); + const Fp8BwdGraphInputs in(cfg); const bool is_delayed_scaling = in.is_delayed_scaling; const bool is_current_scaling = in.is_current_scaling; const bool is_mxfp8 = in.is_mxfp8; @@ -1090,7 +1082,7 @@ void fused_attn_fp8_bwd_impl( dropout_seed, dropout_offset] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built("bwd", *cache_entry); + ensure_plans_built(Backend::FP8, Pass::Bwd, *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1100,7 +1092,7 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } - graph_cache_debug::record_exec("bwd"); + graph_cache_debug::record_exec(Backend::FP8, Pass::Bwd); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1435,7 +1427,7 @@ std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handl graph_cfg.check_for_forward_support = true; try { - const fused_attn::Fp8FwdGraphInputs in = fused_attn::derive_fp8_fwd_graph_inputs(graph_cfg); + const fused_attn::Fp8FwdGraphInputs in(graph_cfg); fused_attn::fp8_fwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception& e) { @@ -1451,7 +1443,7 @@ std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handl graph_cfg.check_for_forward_support = false; try { - const fused_attn::Fp8BwdGraphInputs in = fused_attn::derive_fp8_bwd_graph_inputs(graph_cfg); + const fused_attn::Fp8BwdGraphInputs in(graph_cfg); fused_attn::fp8_bwd_cached_graph(graph_cfg, in, handle); return ""; } catch (const std::exception& e) { diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index dd7f78b11c..7a226c71fb 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -15,6 +15,15 @@ // four, and lives here so it has one definition rather than four copies to keep // in step. // +// The five frontend calls a graph goes through, and which caller pays for each: +// +// on a miss, either caller: +// validate() -> build_operation_graph() -> create_execution_plans(HeurMode_t::A) +// -> check_support() validate_and_check_support() +// the execution path only: +// build_plans() ensure_plans_built(), once per entry, the kernel compilation +// execute() every call, with its variant pack built in a local +// // This header is deliberately not part of utils.h: it needs the cuDNN frontend, // and utils.h is included by translation units (utils.cu) that otherwise do not. // ============================================================================ @@ -22,8 +31,8 @@ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ +#include #include -#include #include #include #include @@ -42,7 +51,7 @@ namespace transformer_engine { namespace fused_attn { // cuDNN's refusal to run a graph, as opposed to a failure to try. The distinction is what makes -// the negative cache in get_or_build_cached_graph() safe: a refusal is a verdict on the +// the negative cache in build_or_get_cached_graph() safe: a refusal is a verdict on the // configuration and reproducible for a given key, so it can be remembered and replayed, whereas // a failure that came from the machine's state at that moment (an allocation that did not fit, a // CUDA error left behind by unrelated work) could well succeed on the next attempt and must not @@ -126,6 +135,16 @@ struct CachedGraph { // relied on for that -- it has held for far longer than the >= 1.25.0 the build requirements // ask for, which is there for unrelated features. // +// What lets one cache serve every thread is an asymmetry between the two objects a call needs. A +// cuDNN handle is per-thread mutable session state: it carries the stream that execute() launches +// on, so each thread holds its own rather than racing to set that on a shared one. A graph and +// its plans are the opposite -- compiled artifacts, built for the properties of a device and +// bound to the device they were finalized against, with nothing in them belonging to the thread +// that did the building. So the cache can be keyed by device and shared by all threads, which is +// why make_cache_key() stamps device_id and nothing thread-shaped. ensure_plans_built() covers +// what that costs at the seam, where the thread that finishes a build is often not the thread +// that started it. +// // Refusals are cached alongside the graphs, under the same keys and the same lock. A support // query for an unsupported configuration is otherwise the most expensive thing this cache sees: // it builds the whole graph, spends the four frontend calls, and throws the result away, and it @@ -139,7 +158,7 @@ struct CachedGraph { // a cache and its lock as two separate objects leaves that ordering to whoever writes the next // one; declaring them here settles it once. // -// Both maps are bounded; see cache_capacity(). `last_used` is what makes the bound an LRU rather +// Both maps are bounded; see kCacheCapacity. `last_used` is what makes the bound an LRU rather // than an arbitrary cull: it is stamped from `clock` on every insertion and every hit, so the // entry with the smallest value is the one that has gone longest without being asked for. The // clock is an ordinary member rather than an atomic because it is only ever touched under @@ -161,42 +180,32 @@ struct GraphCache { std::map unsupported; }; -// The default ceiling on entries in one of the maps of one build site's cache. +// The ceiling on entries in one of the maps of one build site's cache. // // Sized to be out of the way of real work rather than to be tight. A training step reuses a -// handful of configurations, an inference server with bucketed sequence lengths tens of them; -// a few hundred is already far more shape diversity than a model exhibits. What the ceiling is -// for is the case where the key space is effectively unbounded -- a test suite sweeping shapes, -// or a serving workload that keys on something that never repeats -- where an unbounded cache -// is a slow leak of cuDNN graphs and their execution plans for the life of the process. -constexpr size_t kDefaultCacheCapacity = 500; - -// The ceiling in force, from NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES if it is set. 0 means no ceiling, -// which is the escape hatch for a workload that genuinely has thousands of live configurations -// and would rather spend the memory than rebuild. Read once: the limit is a property of the run. -inline size_t cache_capacity() { - static const size_t capacity = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_MAX_ENTRIES"); - if (e == nullptr || e[0] == '\0') return kDefaultCacheCapacity; - const long long v = std::atoll(e); // NOLINT(runtime/int) - return v < 0 ? kDefaultCacheCapacity : static_cast(v); - }(); - return capacity; -} +// handful of configurations and an inference server with bucketed sequence lengths tens of them, +// so a hundred is already more shape diversity than a model exhibits. What the ceiling is for is +// the case where the key space is effectively unbounded -- a test suite sweeping shapes, or a +// serving workload that keys on something that never repeats -- where an unbounded cache is a +// slow leak of cuDNN graphs and their execution plans for the life of the process. +// +// Hard-coded rather than configurable, because nothing has yet needed a different number: the +// workloads that fit under it never notice the ceiling, and the ones that do not are better +// served by rebuilding a graph than by holding thousands. An environment variable can come back +// if a workload turns up that wants to trade the memory for the rebuilds. +constexpr size_t kCacheCapacity = 100; // Make room in `entries` for one more, by dropping the least recently used until there is. // Call under the cache's lock. // -// Evicting a graph does not invalidate one that is in use. get_or_build_cached_graph() hands +// Evicting a graph does not invalidate one that is in use. build_or_get_cached_graph() hands // back a shared_ptr, so a thread that is executing an entry holds it alive regardless of what // the map does; erasing here drops the cache's reference and nothing else. The scan is linear, -// but it runs only when the cache is full, and comparing a few hundred integers is nothing -// beside the graph build it is making room for. +// but it runs only when the cache is full, and comparing a hundred integers is nothing beside +// the graph build it is making room for. template void evict_to_fit(Map &entries) { - const size_t capacity = cache_capacity(); - if (capacity == 0) return; - while (entries.size() >= capacity) { + while (entries.size() >= kCacheCapacity) { auto oldest = entries.begin(); for (auto it = entries.begin(); it != entries.end(); ++it) { if (it->second.last_used < oldest->second.last_used) oldest = it; @@ -207,8 +216,8 @@ void evict_to_fit(Map &entries) { // Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: // validate, build_operation_graph, create_execution_plans, check_support. The sequence is -// identical for both passes and both backends, so it is defined once here; `pass` only selects -// which set of stage timers the calls are attributed to. +// identical for both passes and both backends, so it is defined once here; `backend` and `pass` +// only name the build site whose stage timers the calls are attributed to. // // Support is reported by throwing rather than by a return value. NVTE_CHECK_CUDNN_FE raises // an exception carrying cuDNN's own explanation of the rejection, and that text is what the @@ -228,11 +237,12 @@ void evict_to_fit(Map &entries) { // build_plans() and execute() sit outside this function entirely: they commit real resources, and // build_plans() belongs to whoever executes the graph, once, the first time it is needed. See // CachedGraph. -inline void validate_and_check_support(const char *pass, cudnn_frontend::graph::Graph &graph, - cudnnHandle_t handle) { +inline void validate_and_check_support(graph_cache_debug::Backend backend, + graph_cache_debug::Pass pass, + cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { cudnn_frontend::error_t error; - graph_cache_debug::timer(pass, stage, [&] { error = call(); }); + graph_cache_debug::timer(backend, pass, stage, [&] { error = call(); }); if (error.is_good()) return; // cuDNN normally explains itself; fall back to the call's name so that a refusal can never // arrive as an empty string, which the is_supported_* helpers would read as an endorsement. @@ -272,13 +282,32 @@ inline void validate_and_check_support(const char *pass, cudnn_frontend::graph:: // holding the lock across it would serialize builds of unrelated keys, so two threads racing // on the same key may both build. That is a wasted build, not a correctness problem: the // loser drops its own graph and takes the winner's, so every caller of a given key gets one -// shared entry and the once-flag inside it still governs the plan build. The wasted build is -// visible in diagnostics as a BUILD with no matching MISS of its own. The same race on a -// refused key is equally harmless, both threads storing the same reason. +// shared entry and the once-flag inside it still governs the plan build. Both threads record +// their own lookup, so the wasted build shows up in diagnostics as two MISS lines carrying the +// same key and a build_graph count above the number of distinct keys, rather than as anything +// missing. The same race on a refused key is equally harmless, both threads storing the same +// reason. +// +// lock cache.mutex +// supported[key]? found -> last_used = ++clock, copy the shared_ptr +// unsupported[key]? found -> last_used = ++clock, copy the reason +// unlock +// record_cache_lookup(HIT | UNSUPPORTED | MISS) +// +// HIT -> return the entry +// UNSUPPORTED -> throw UnsupportedGraph(the remembered reason) +// MISS -> build() outside the lock, so builds of unrelated +// validate_and_check_support() keys proceed concurrently +// ok -> lock, evict_to_fit(supported), insert stamped ++clock, unlock, +// return the inserted entry, which on a lost race is the winner's +// verdict -> lock, evict_to_fit(unsupported), insert the reason, unlock, +// rethrow +// other -> NVTE_ERROR: nothing remembered, retried when the key returns template -std::shared_ptr> get_or_build_cached_graph( - GraphCache &cache, const FusedAttnConfig &key, const char *pass, - cudnnHandle_t handle, BuildFn &&build) { +std::shared_ptr> build_or_get_cached_graph( + GraphCache &cache, const FusedAttnConfig &key, + graph_cache_debug::Backend backend, graph_cache_debug::Pass pass, cudnnHandle_t handle, + BuildFn &&build) { using Entry = CachedGraph; using Slot = typename GraphCache::Slot; using Refusal = typename GraphCache::Refusal; @@ -312,7 +341,7 @@ std::shared_ptr> get_or_build_cached_graph( // querying other keys. The counters are exact, but two lookups that raced on the lock can be // recorded in the opposite order, so read a level-2 trace as the set of lookups that happened // rather than as the sequence they happened in. - graph_cache_debug::record_cache_lookup(pass, outcome, key); + graph_cache_debug::record_cache_lookup(backend, pass, outcome, key); if (cached != nullptr) return cached; // Raised rather than returned so that a replayed refusal is the same event as a fresh one: @@ -326,17 +355,17 @@ std::shared_ptr> get_or_build_cached_graph( // Every site's tensor tuple leads with its graph, which is the one thing all four have in // common and the only element this needs. A tuple that stopped leading with it would fail to // compile here rather than quietly validate the wrong object. - validate_and_check_support(pass, *std::get<0>(entry->tensors), handle); + validate_and_check_support(backend, pass, *std::get<0>(entry->tensors), handle); } catch (const UnsupportedGraph &e) { { std::lock_guard lock(cache.mutex); evict_to_fit(cache.unsupported); cache.unsupported.insert({key, Refusal{e.what(), ++cache.clock}}); } - graph_cache_debug::record_unsupported(pass); + graph_cache_debug::record_unsupported(backend, pass); throw; } - graph_cache_debug::record_build(pass); + graph_cache_debug::record_graph_built(backend, pass); { std::lock_guard lock(cache.mutex); evict_to_fit(cache.supported); @@ -347,20 +376,45 @@ std::shared_ptr> get_or_build_cached_graph( } } -// Runs the plan build that get_or_build_cached_graph() left undone, once per entry. +// Runs the plan build that build_or_get_cached_graph() left undone, once per entry. // // Call this only when the graph is about to be executed, which is why it is a separate step // rather than the tail of the lookup: a support query builds entries that nothing ever runs, and // kernel compilation is the most expensive of the five frontend calls, so a query that paid for // it would be paying for nothing. See CachedGraph for why the flag lives inside the entry and // what a throw here leaves behind. +// +// Splitting the build in two means the thread that finishes it is often not the thread that +// started it -- a sizing call on one thread caches the graph, and an autograd thread is the first +// to need it to run. Four facts make that safe, and only the first is visible here. +// +// build_plans() takes no handle. The overload that accepts one ignores it -- its body is +// `(void)handle;` -- and the build works from the operation graph descriptor and the device +// properties instead, which is how deviceless ahead-of-time compilation builds plans with no +// handle at all. Unlike the plan sharing described on GraphCache, this does lean on the >= 1.25.0 +// frontend the build requires: it is where the handle-free overload arrived. Calling it means a +// plan build cannot reach for the handle of a thread that has since exited. +// +// The handle from the build does outlive the build, held by the operation graph descriptor that +// build_operation_graph(handle) finalized against it. It stays a valid object only because TE +// never destroys cuDNN handles: cudnnExecutionPlanManager leaves HandleManager's Destroy +// parameter at its nullptr default, so handles leak by design, one per thread per device. +// +// That descriptor was finalized for the device of the handle that built it, which is why the +// cache key carries device_id (see FusedAttnConfig::make_cache_key). Without it a thread could +// build plans, and compile kernels, from a descriptor belonging to another device. +// +// Execution stays clear of all of it: execute() is called with the running thread's own handle, +// so a handle is never used by two threads at once, which is what cuDNN asks in return for +// letting them share the plan. template -void ensure_plans_built(const char *pass, CachedGraph &entry) { +void ensure_plans_built(graph_cache_debug::Backend backend, graph_cache_debug::Pass pass, + CachedGraph &entry) { std::call_once(entry.plans_built, [&] { cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.tensors); - graph_cache_debug::timer(pass, graph_cache_debug::BuildStage::BuildPlans, + graph_cache_debug::timer(backend, pass, graph_cache_debug::BuildStage::BuildPlans, [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); - graph_cache_debug::record_plans_built(pass); + graph_cache_debug::record_plans_built(backend, pass); }); } diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 84c51e3cd2..4222843ce9 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -8,55 +8,72 @@ // Fused-attention graph cache diagnostics. // // Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG. Two verbosity levels: -// =1 : low volume. Cache event counters, a BUILD and a PLANS line per build, an -// UNSUP line per configuration cuDNN refuses, and the end-of-run SUMMARY -// (aggregate + per-thread) and stage timings. This is enough to diagnose -// redundant rebuilds and profile build cost. -// =2 : high volume (trace). Additionally emits a per-lookup HIT/MISS/NOSUP line -// with the full shorthand cache key and a per-execution EXEC line. Use only -// when you need to see *which* shapes are hitting/missing -- these fire on -// every cache lookup and execution, so at suite scale they add I/O and -// serialize threads on the stderr lock. No timed region writes to stderr, so -// the stage timings stay sound, but they are measured under more contention -// than at level 1 and read a little high. +// =1 (events) : low volume. Cache event counters, a BUILD_GRAPH and a BUILD_PLANS line +// per build, an UNSUPPORTED line per configuration cuDNN refuses, and the +// end-of-run SUMMARY (per backend and per thread, plus a row across the +// backends when a run used more than one) with stage timings. Each of these +// fires once per distinct cache key, which is what keeps the volume low, and +// is enough to diagnose redundant rebuilds and profile build cost. +// =2 (trace) : high volume. Additionally emits a per-lookup HIT/MISS/UNSUPPORTED line +// with the full shorthand cache key and a per-execution EXEC line. Use only +// when you need to see *which* shapes are hitting/missing -- these fire on +// every cache lookup and execution, so at suite scale they add I/O and +// serialize threads on the stderr lock. No timed region writes to stderr, so +// the stage timings stay sound, but they are measured under more contention +// than at level 1 and read a little high. // -// NOSUP is a hit on the negative cache: a key cuDNN has already refused, answered -// from the stored refusal instead of by building the graph again. +// Every line names the build site behind it, "f16" or "fp8" followed by the pass, and the +// counters it carries belong to that backend alone -- the two keep separate columns, so a +// process that drives both can still say which of them built what. Every event name is also +// the counter column it increments, so a line and the totals beside it read with one +// vocabulary. UNSUPPORTED names both a level-1 event and a level-2 lookup outcome, which are +// the two halves of one story: the event records the refusal cuDNN just handed back, and the +// lookup line is a later query answered from that stored refusal instead of by building the +// graph again. Tell them apart by the line shape -- the event line carries counters, the +// lookup line carries the cache key. // // An optional ":" suffix picks which processes emit, defaulting to rank 0 // so that output does not scale with the world size: "1:all" for every rank, // "2:0,3" for a specific set. See `rank_selected` for when overriding pays off. // // Level 1 on one training step of a supported configuration. Every line begins with -// "[FUSED-ATTN-CACHE] pid=[ rank=] | ", elided below, and carries the running -// totals, of which only the pass being reported is shown (the counters are printed -// right-aligned in a fixed width, dropped here): +// "[FUSED-ATTN-CACHE] rank= | ", or with just "[FUSED-ATTN-CACHE] " when the launcher +// exports no rank (see `rank_tag`), elided below, and carries the running totals, of which +// only the pass being reported is shown (the counters are printed right-aligned in a fixed +// width, and are abbreviated here): // -// THREAD | tid=0 dev=0 os_tid=1234 -// fwd BUILD | tid=0 dev=0 | fwd miss=1, hit=0, built=1, unsup=0, plans=0, exec=0 | bwd ... -// bwd BUILD | tid=0 dev=0 | fwd ... | bwd miss=1, hit=0, built=1, unsup=0, plans=0, ... -// fwd PLANS | tid=0 dev=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=0 | bwd ... +// f16 fwd BUILD_GRAPH | tid=0 dev=0 | fwd hit_supported=0, miss=1, build_graph=1, ... +// f16 bwd BUILD_GRAPH | tid=0 dev=0 | fwd ... | bwd hit_supported=0, miss=1, ... +// f16 fwd BUILD_PLANS | tid=0 dev=0 | fwd hit_supported=1, miss=1, build_plans=1, ... // ===== summary begin ===== -// SUMMARY-TID | tid=0 dev=0 | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... -// SUMMARY | tid=all dev=all | fwd miss=1, hit=1, built=1, unsup=0, plans=1, exec=1 | bwd ... -// fwd check_support | calls=1 | time= 42.135 ms/call +// f16 SUMMARY-TID | tid=0 dev=0 | fwd hit_supported=5, miss=1, build_graph=1, ... +// f16 SUMMARY-TID | tid=1 dev=0 | fwd ... | bwd hit_supported=4, build_plans=1, ... +// f16 SUMMARY | tid=all dev=all | fwd hit_supported=5, miss=1, build_graph=1, ... +// f16 fwd check_support | calls=1 | time= 0.031 ms/call +// f16 fwd build_plans | calls=1 | time= 262.104 ms/call // ===== summary end ===== // +// The two thread rows are what a PyTorch step really looks like: the forward, and the support +// probe for the backward, run on the main thread, while the backward itself runs on the +// autograd thread and finds the graph that probe left behind. Neither row satisfies +// `build_graph >= build_plans` by itself -- tid=1 compiled the plans of a graph tid=0 built -- +// so read the identities off the totals rows rather than the per-thread ones. +// // The device column matters as soon as one process drives more than one -- device_id is part of // the cache key, so the same shape on two devices is two entries, and a build count that looks -// doubled is explained by reading which device each BUILD came from. +// doubled is explained by reading which device each BUILD_GRAPH came from. // -// Two forward lookups against one build is the shape of a healthy run: the support -// query missed and built, and the execution that followed hit the entry the query left -// behind. `built=1, plans=1` says that graph went on to be executed; `built` above -// `plans` counts graphs built for a query and never run. A refused configuration reads -// `miss=1, unsup=1, built=0` instead, and stays at one refusal however many times it is -// queried. +// A support query misses and builds, and every later lookup of that key is a hit_supported -- +// including the workspace-sizing call that precedes each execution -- so the hit columns climb +// faster than exec. `build_graph=1, build_plans=1` says that graph went on to be executed; +// `build_graph` above `build_plans` counts graphs built for a query and never run. A refused +// configuration reads `miss=1, unsupported=1, build_graph=0` instead, and stays at one refusal +// however many times it is queried: the repeat queries land in hit_unsupported. // // Level 2 adds one line per lookup and per execution, with the key that decided it: // -// fwd MISS | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... -// fwd HIT | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... +// f16 fwd MISS | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... +// f16 fwd HIT | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... // // where diffing two MISS lines names the fields that cost the extra build. // ============================================================================ @@ -64,9 +81,6 @@ #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 @@ -88,6 +102,19 @@ namespace transformer_engine { namespace fused_attn { namespace graph_cache_debug { +// Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=events, +// 2=trace). Single read at startup, cached; when unset every call site pays one +// cached-flag check and nothing else. +inline int debug_level() { + static const int lvl = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; + const int v = std::atoi(e); // stops at the optional ":" suffix + return v > 0 ? v : 1; // any non-empty, non-"0" value enables at least level 1 + }(); + return lvl; +} + // Rank of this process as reported by the launcher, or -1 when there is no // launcher (a single-process run). First variable that is set wins. inline int launcher_rank() { @@ -101,18 +128,6 @@ inline int launcher_rank() { return rank; } -// Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=default, -// 2=trace). Single read at startup, cached. Negligible overhead when unset. -inline int debug_level() { - static const int lvl = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; - const int v = std::atoi(e); // stops at the optional ":" suffix - return v > 0 ? v : 1; // any non-empty, non-"0" value enables at least level 1 - }(); - return lvl; -} - // Whether this process emits diagnostics. Every rank writes to the same stderr, // so emitting from all of them multiplies the volume by the world size -- and // under data/tensor parallelism the ranks are running identical shapes, so the @@ -146,73 +161,210 @@ inline bool rank_selected() { // Diagnostics are on at level >= 1, and only for the selected ranks. Unselected // ranks skip the counters too, so they pay nothing beyond this check. -inline bool enabled() { return debug_level() >= 1 && rank_selected(); } +// +// Cached in its own flag rather than recomputed from the two above, so that this -- the check +// every call site makes, on the per-lookup path included -- reads one initialized-once static +// instead of two. Both inputs are fixed for the life of the process, so there is nothing to +// recompute; `rank_selected` is still only reached when the level says diagnostics are on. +inline bool enabled() { + static const bool on = debug_level() >= 1 && rank_selected(); + return on; +} // Per-lookup / per-exec trace lines are gated behind level >= 2. inline bool trace_enabled() { return debug_level() >= 2; } -// Identifies the emitting process. Distributed PyTorch runs one process per rank -// and they all share this stderr, so without this every line would be ambiguous -// (thread ids restart at 0 in each process). Rank comes from the launcher, if any. -inline const std::string &process_tag() { +// Names the emitting rank. Distributed runs put one process per rank on the same stderr, so +// without this the ranks' lines would be indistinguishable. A run whose launcher exports no +// rank is left untagged rather than falling back to a pid: an OS-level identifier is only +// useful for correlating against a profiler or another process, which these logs are not for. +// The tag carries its own trailing separator, so the untagged case prints no empty column. +inline const std::string &rank_tag() { static const std::string *tag = [] { - auto *s = new std::string("pid=" + std::to_string(static_cast(::getpid()))); - if (launcher_rank() >= 0) *s += " rank=" + std::to_string(launcher_rank()); - return s; + const int rank = launcher_rank(); + if (rank < 0) return new std::string(); + return new std::string("rank=" + std::to_string(rank) + " | "); }(); return *tag; } -// More readable, shorter thread IDs (0, 1, 2, ...). These are assignment order, -// not identity: tid=0 is whichever thread touched this cache first. The one-shot -// THREAD line below maps them to OS thread ids for correlating with nsys/gdb. +// More readable, shorter thread IDs (0, 1, 2, ...). These are assignment order, not identity: +// tid=0 is whichever thread touched this cache first, and the number means nothing outside this +// process. It exists to attribute the per-thread SUMMARY rows, not to be matched against +// anything external. 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; } -// OS-level thread id, as reported by nsys/gdb/`top -H`. Via syscall rather than -// gettid() so this does not require glibc >= 2.30. -inline int64_t os_thread_id() { return static_cast(::syscall(SYS_gettid)); } - // 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 graph built and cached in response to a cache miss. Built only as far as -// check_support(), which is all a support probe needs. -// - PLANS: a cached graph finished with build_plans(), the kernel compilation that the -// BUILD above deferred. At most one per BUILD, on the first execution of that -// graph, so BUILD minus PLANS is how many graphs were built for a support probe -// and never used to run anything. -// - 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 +// The build site an event came from: f16 or fp8, forward or backward. Every recorder names +// both halves, because the counters are kept per site rather than per pass. One process can +// drive both backends, and adding f16's builds into the same column as fp8's would leave such +// a run unable to say which of them paid for what. +// +// Backend::F16 is the arbitrary-seqlen f16 backend; the max512 one keeps no graph cache and so +// has nothing to report here. Naming the site with a pair of enums rather than with the +// "fwd"/"bwd" strings this used to take is also what turns a mistake at a call site into a +// compile error instead of an event silently counted against the wrong column. +// ============================================================================ + +enum class Backend { F16, FP8 }; +enum class Pass { Fwd, Bwd }; + +inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } +inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } + +// Backend major, pass minor, so that the two passes of one backend are adjacent -- which is how +// the counter lines and the summary rows present them, one backend at a time. +constexpr size_t kSiteCount = 4; +inline constexpr size_t site_index(Backend b, Pass p) { + return (b == Backend::F16 ? 0u : 2u) + (p == Pass::Fwd ? 0u : 1u); +} + +// ============================================================================ +// Cache event counters, one block per build site. Each name is both the event tag on the line +// that records it and the column carrying its running total: +// - build_graph: a graph built and cached in response to a cache miss. Built only as far +// as check_support(), which is all a support probe needs. +// - build_plans: a cached graph finished with build_plans(), the kernel compilation that +// build_graph deferred. At most one per build_graph, and paid by the first +// execution of that graph rather than by the probe that built it. +// - unsupported: a configuration cuDNN refused, now remembered as a negative cache entry. +// The other way a miss can end. Counted once per refusal recorded, which is +// normally once per distinct refused key; later queries for it are +// hit_unsupported. +// - exec: a graph execution call with valid runtime tensors +// - hit_supported: a lookup answered from the graph map. May not lead to an exec: it can be +// a backend availability check, or the workspace-sizing call of +// nvte_fused_attn_fwd/bwd, which has no runtime tensors to run with. +// - hit_unsupported: a lookup answered from the refusal map -- a key cuDNN has already +// refused, replayed instead of rebuilt. Both hit columns are named for the +// map that answered them, and `unsupported` above counts the refusals +// themselves rather than the queries that replay them. +// - miss: a lookup neither map answered; triggers a graph build +// +// Identities. These hold by construction, so a violation is a bug in the cache or in the +// counting rather than something the workload did: +// - hit_supported + hit_unsupported + miss = every lookup, one recorded per entry into +// build_or_get_cached_graph, which makes it the denominator for everything below. +// - miss = build_graph + unsupported. A shortfall in either means a build ended in +// something cuDNN did not state as a verdict on the graph. +// - build_graph >= build_plans, the gap being graphs a probe built that nothing has run. +// Eviction grows both rather than closing it: a rebuilt key gets a fresh once_flag. +// - exec > 0 implies build_plans > 0, every site calling ensure_plans_built ahead of the +// workspace-sizing return, which is itself ahead of record_exec. The same ordering read +// backwards: a workspace-sizing call pays build_plans and never exec. +// - hit_unsupported > 0 implies unsupported > 0, a refusal being replayable only once some +// earlier call has recorded it. +// - The two build identities are properties of the totals rows, not of one SUMMARY-TID row: +// the thread that builds a graph need not be the thread that compiles its plans, and a +// PyTorch step splits exactly that way across the autograd thread. +// - A backend's SUMMARY-TID rows sum column by column to its SUMMARY row, and the +// per-backend rows to the all-backends one. +// - A lost build race disturbs none of the above: the loser records its own miss and its own +// build_graph, so both sides of miss = build_graph + unsupported move together, and the +// entry's once_flag still permits only one build_plans. What a race does break is reading +// build_graph as the number of graphs cached, two builds being able to stand behind one +// entry; the same goes for unsupported and the number of keys cuDNN has refused. +// - In the stage timing rows, calls only fall along the sequence validate >= +// build_operation_graph >= create_execution_plans >= check_support, each drop being the +// builds that ended at the stage before -- which localizes where cuDNN refuses, rather +// than only how long refusing took. +// - The build_plans timing row can show more calls than the build_plans column counts, the +// difference being plan builds that threw: the timer records while unwinding, the counter +// only after the call returns. +// +// Signatures. Workload-dependent, so these are read rather than asserted: +// - After warmup only hit_supported and exec should move. A build_graph late in a run means +// something varies per step that need not. +// - Several hit_supported per exec is normal, since backend selection, workspace sizing and +// execution all look the same key up; what matters is that the ratio stays flat. +// - exec / build_graph is the amortization figure, how many executions each built graph +// served, and a lower bound at that, since a race or an eviction adds a build without +// adding a graph. Single digits after a long run means the cache is not earning its keep. +// - hit_unsupported climbing while unsupported stays at one is the negative cache doing its +// job. It also says this site never runs fused, which makes it the column to reach for +// when attention is slower than expected and nothing raised an error. +// - build_graph or unsupported past kCacheCapacity suggests that map has evicted. A hint +// rather than an identity: a lost build race counts twice against one key. +// - Two MISS lines carrying the same key, with build_graph above the number of distinct keys, +// is that lost race. It is wasted work rather than a bug, and worth chasing only if it +// repeats, which would mean threads are arriving on cold keys together every step. +// - A level-2 trace is the set of lookups that happened, not the order they happened in: the +// line is written after the cache lock is dropped, so two threads that raced for it can +// print in the opposite order. // ============================================================================ struct EventCounters { - std::atomic built{0}; - std::atomic plans{0}; + std::atomic build_graph{0}; + std::atomic build_plans{0}; std::atomic exec{0}; - std::atomic hit{0}; + std::atomic hit_supported{0}; + std::atomic hit_unsupported{0}; std::atomic miss{0}; - std::atomic unsup{0}; + std::atomic unsupported{0}; +}; + +inline EventCounters &counters(Backend b, Pass p) { + static std::array table{}; + return table[site_index(b, p)]; +} + +// One counter block read out into plain values. The summary sums blocks to get its per-backend +// and all-backends rows, atomics cannot be summed, and this is where the reading happens; it +// also keeps the loads out of the formatting. The columns are not read as one indivisible +// operation, which nothing here wants: the summary runs at exit, after the threads that wrote +// them are done, and an event line is a snapshot of a moving count by nature. +struct CounterSnapshot { + uint64_t build_graph = 0; + uint64_t build_plans = 0; + uint64_t exec = 0; + uint64_t hit_supported = 0; + uint64_t hit_unsupported = 0; + uint64_t miss = 0; + uint64_t unsupported = 0; + + CounterSnapshot &operator+=(const CounterSnapshot &other) { + build_graph += other.build_graph; + build_plans += other.build_plans; + exec += other.exec; + hit_supported += other.hit_supported; + hit_unsupported += other.hit_unsupported; + miss += other.miss; + unsupported += other.unsupported; + return *this; + } + + // Whether this block saw nothing at all, which is what lets the summary leave out the rows + // for a backend the run never used rather than printing zeros for it. + bool empty() const { + return (build_graph | build_plans | exec | hit_supported | hit_unsupported | miss | + unsupported) == 0; + } }; -inline EventCounters &counters(bool is_fwd) { - static EventCounters fwd; - static EventCounters bwd; - return is_fwd ? fwd : bwd; +inline CounterSnapshot snapshot(const EventCounters &c) { + CounterSnapshot s; + s.build_graph = c.build_graph.load(std::memory_order_relaxed); + s.build_plans = c.build_plans.load(std::memory_order_relaxed); + s.exec = c.exec.load(std::memory_order_relaxed); + s.hit_supported = c.hit_supported.load(std::memory_order_relaxed); + s.hit_unsupported = c.hit_unsupported.load(std::memory_order_relaxed); + s.miss = c.miss.load(std::memory_order_relaxed); + s.unsupported = c.unsupported.load(std::memory_order_relaxed); + return s; } -// Per-thread counters, so the summary can break down build/exec/hit/miss by -// thread. In the single-process context-parallel case each device is driven by -// its own thread, so this reveals which thread built/executed what. +// Per-thread counters, one block per build site, so the summary can break down every column by +// thread and backend. In the single-process context-parallel case each device is driven by its +// own thread, so this reveals which thread built and executed what; under PyTorch it also +// separates the main thread from the autograd thread that runs the backward. // // `device` is the device this thread last drove, restamped on every event. The event lines print // the live current device, which is exact; this exists for the SUMMARY-TID rows, which are @@ -222,8 +374,7 @@ inline EventCounters &counters(bool is_fwd) { struct ThreadCounters { unsigned tid = 0; std::atomic device{-1}; - EventCounters fwd; - EventCounters bwd; + std::array sites; }; // The registry and its mutex are heap-allocated and deliberately never freed. @@ -254,130 +405,146 @@ inline ThreadCounters &thread_counters() { // Stamped here as well as on every event, so that a thread which only ever hits the cache -- // and so never reaches print_counters() at level 1 -- still names a device in the summary // rather than reporting the -1 it was constructed with. - const int device = cuda::current_device(); - p->device.store(device, std::memory_order_relaxed); + p->device.store(cuda::current_device(), std::memory_order_relaxed); { std::lock_guard lock(thread_registry_mutex()); thread_registry().push_back(p); } - // One line per thread, mapping the short id to something nsys/gdb can match. - std::fprintf(stderr, - "[FUSED-ATTN-CACHE] %s | THREAD | tid=%-3u dev=%-3d os_tid=%" PRId64 "\n", - process_tag().c_str(), p->tid, device, os_thread_id()); - std::fflush(stderr); return p; }(); return *tc; } -inline EventCounters &thread_counters(bool is_fwd) { - ThreadCounters &tc = thread_counters(); - return is_fwd ? tc.fwd : tc.bwd; +inline EventCounters &thread_counters(Backend b, Pass p) { + return thread_counters().sites[site_index(b, p)]; } -// Format one counter block (aggregate or a single thread's) as one line. -// `tid_field` is the whole thread column, e.g. "tid=3"; the aggregate row passes -// "tid=all" so that it cannot be misread as thread 0's row. `dev_field` is the device column and -// works the same way, "dev=all" on the aggregate row -- the counters there are summed across -// whatever devices the process drove, so naming one of them would be a lie. +// Format one pair of counter blocks -- the two passes of a single backend -- as one line. +// `label` is the event or summary tag, and names the backend whenever the line speaks for one. +// `tid_field` is the whole thread column, e.g. "tid=3"; the totals rows pass "tid=all" so that +// they cannot be misread as thread 0's row. `dev_field` is the device column and works the same +// way, "dev=all" on a totals row -- those counters are summed across whatever devices the +// process drove, so naming one of them would be a lie. // -// The columns are meant to be read against two identities. Every lookup lands in exactly one of -// miss and hit, and every miss ends in exactly one of built and unsup -- so `miss = built + unsup` -// and a shortfall in either means a build died of something other than a refusal. `built >= plans` -// always, the difference being graphs that a support query built and nothing has yet run. -inline std::string format_counter_line(const char *event, const char *tid_field, - const char *dev_field, const EventCounters &f, - const EventCounters &b) { +// What the columns mean, the identities they can be asserted against and the ratios worth +// reading are all with the counter definitions above. +inline std::string format_counter_line(const char *label, const char *tid_field, + const char *dev_field, const CounterSnapshot &f, + const CounterSnapshot &b) { char buf[768]; std::snprintf(buf, sizeof(buf), - "[FUSED-ATTN-CACHE] %s | %-11s | %-7s %-7s | fwd miss=%4" PRIu64 ", hit=%4" PRIu64 - ", built=%4" PRIu64 ", unsup=%4" PRIu64 ", plans=%4" PRIu64 ", exec=%4" PRIu64 - " | bwd miss=%4" PRIu64 ", hit=%4" PRIu64 ", built=%4" PRIu64 ", unsup=%4" PRIu64 - ", plans=%4" PRIu64 ", exec=%4" PRIu64 "\n", - process_tag().c_str(), event, tid_field, dev_field, - f.miss.load(std::memory_order_relaxed), f.hit.load(std::memory_order_relaxed), - f.built.load(std::memory_order_relaxed), f.unsup.load(std::memory_order_relaxed), - f.plans.load(std::memory_order_relaxed), f.exec.load(std::memory_order_relaxed), - b.miss.load(std::memory_order_relaxed), b.hit.load(std::memory_order_relaxed), - b.built.load(std::memory_order_relaxed), b.unsup.load(std::memory_order_relaxed), - b.plans.load(std::memory_order_relaxed), b.exec.load(std::memory_order_relaxed)); + "[FUSED-ATTN-CACHE] %s%-19s | %-7s %-7s | fwd hit_supported=%4" PRIu64 + ", hit_unsupported=%4" PRIu64 ", miss=%4" PRIu64 ", build_graph=%4" PRIu64 + ", unsupported=%4" PRIu64 ", build_plans=%4" PRIu64 ", exec=%4" PRIu64 + " | bwd hit_supported=%4" PRIu64 ", hit_unsupported=%4" PRIu64 ", miss=%4" PRIu64 + ", build_graph=%4" PRIu64 ", unsupported=%4" PRIu64 ", build_plans=%4" PRIu64 + ", exec=%4" PRIu64 "\n", + rank_tag().c_str(), label, tid_field, dev_field, f.hit_supported, f.hit_unsupported, + f.miss, f.build_graph, f.unsupported, f.build_plans, f.exec, b.hit_supported, + b.hit_unsupported, b.miss, b.build_graph, b.unsupported, b.build_plans, b.exec); return std::string(buf); } -inline void print_counter_block(const char *event, const char *tid_field, const char *dev_field, - const EventCounters &f, const EventCounters &b) { - const std::string line = format_counter_line(event, tid_field, dev_field, f, b); +inline void print_counter_block(const char *label, const char *tid_field, const char *dev_field, + const CounterSnapshot &f, const CounterSnapshot &b) { + const std::string line = format_counter_line(label, tid_field, dev_field, f, b); std::fputs(line.c_str(), stderr); std::fflush(stderr); } -// One event line, from the thread the event happened on. The device is read live rather than -// remembered, so it is the device this event was actually issued against, and is recorded on the -// thread's block on the way past for the benefit of the exit summary. -inline void print_counters(const char *event) { +// One event line, from the thread the event happened on, carrying the running totals of the +// backend that raised it. The device is read live rather than remembered, so it is the device +// this event was actually issued against, and is recorded on the thread's block on the way past +// for the benefit of the exit summary. +inline void print_counters(Backend b, Pass p, const char *event) { const int device = cuda::current_device(); thread_counters().device.store(device, std::memory_order_relaxed); + char label[32]; char tid_field[16]; char dev_field[16]; + std::snprintf(label, sizeof(label), "%s %s %s", backend_name(b), pass_name(p), event); std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); - print_counter_block(event, tid_field, dev_field, counters(/*is_fwd=*/true), - counters(/*is_fwd=*/false)); + print_counter_block(label, tid_field, dev_field, snapshot(counters(b, Pass::Fwd)), + snapshot(counters(b, Pass::Bwd))); } // A graph built through check_support() and cached. Call after the build, from the miss // path that performed it. -inline void record_build(const char *pass) { +inline void record_graph_built(Backend b, Pass p) { 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); - thread_counters(is_fwd).built.fetch_add(1, std::memory_order_relaxed); - print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD"); + counters(b, p).build_graph.fetch_add(1, std::memory_order_relaxed); + thread_counters(b, p).build_graph.fetch_add(1, std::memory_order_relaxed); + print_counters(b, p, "BUILD_GRAPH"); } -// The build_plans() a BUILD deferred, now completed. Call from inside the std::call_once +// The build_plans() a build_graph deferred, now completed. Call from inside the std::call_once // that runs it, after the call returns rather than before: build_plans() throws without // setting the once_flag, leaving a later execution to retry it, so counting on the way out -// keeps this a count of graphs that reached a runnable state. Like BUILD this fires once +// keeps this a count of graphs that reached a runnable state. Like build_graph this fires once // per distinct cache key, so it stays on the level-1 path. -inline void record_plans_built(const char *pass) { +inline void record_plans_built(Backend b, Pass p) { if (!enabled()) return; register_summary_once(); - const bool is_fwd = std::strcmp(pass, "fwd") == 0; - counters(is_fwd).plans.fetch_add(1, std::memory_order_relaxed); - thread_counters(is_fwd).plans.fetch_add(1, std::memory_order_relaxed); - print_counters(is_fwd ? "fwd PLANS" : "bwd PLANS"); + counters(b, p).build_plans.fetch_add(1, std::memory_order_relaxed); + thread_counters(b, p).build_plans.fetch_add(1, std::memory_order_relaxed); + print_counters(b, p, "BUILD_PLANS"); } // A build that cuDNN refused, now remembered as a negative cache entry. Call from the miss path -// that attempted it, in place of record_build(): a refusal and a build are the two ways a miss -// can end, and counting both keeps `miss = built + unsup` true. Fires once per distinct refused -// key -- the second query for that key is a hit -- so it stays on the level-1 path. -inline void record_unsupported(const char *pass) { +// that attempted it, in place of record_graph_built(): a refusal and a build are the two ways a +// miss can end, and counting both keeps `miss = build_graph + unsupported` true. Fires once per +// refused key -- later queries for it land in hit_unsupported -- so it stays on the level-1 path. +inline void record_unsupported(Backend b, Pass p) { if (!enabled()) return; register_summary_once(); - const bool is_fwd = std::strcmp(pass, "fwd") == 0; - counters(is_fwd).unsup.fetch_add(1, std::memory_order_relaxed); - thread_counters(is_fwd).unsup.fetch_add(1, std::memory_order_relaxed); - print_counters(is_fwd ? "fwd UNSUP" : "bwd UNSUP"); + counters(b, p).unsupported.fetch_add(1, std::memory_order_relaxed); + thread_counters(b, p).unsupported.fetch_add(1, std::memory_order_relaxed); + print_counters(b, p, "UNSUPPORTED"); } -inline void record_exec(const char *pass) { +inline void record_exec(Backend b, Pass p) { 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); - thread_counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); + counters(b, p).exec.fetch_add(1, std::memory_order_relaxed); + thread_counters(b, p).exec.fetch_add(1, std::memory_order_relaxed); // The per-exec line fires on every execution; keep it out of the level-1 path. if (!trace_enabled()) return; - print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC"); + print_counters(b, p, "EXEC"); } // What a lookup found. Unsupported is the negative-cache case: a key whose graph cuDNN has // already refused, so the answer is a remembered refusal rather than a graph. enum class LookupResult { Miss, Hit, Unsupported }; +// The column a lookup lands in, which is the cache map that answered it. Written as a switch +// with no default so that adding an outcome fails to compile here rather than being silently +// counted as a miss. +inline std::atomic &lookup_column(EventCounters &c, LookupResult result) { + switch (result) { + case LookupResult::Hit: + return c.hit_supported; + case LookupResult::Unsupported: + return c.hit_unsupported; + case LookupResult::Miss: + break; + } + return c.miss; +} + +inline const char *lookup_name(LookupResult result) { + switch (result) { + case LookupResult::Hit: + return "HIT"; + case LookupResult::Unsupported: + return "UNSUPPORTED"; + case LookupResult::Miss: + break; + } + return "MISS"; +} + // `key` is the normalized cache key -- make_cache_key()'s output, the exact value the // lookup was performed with -- not the execution config it was derived from. That is // deliberate: HIT/MISS is decided by comparing keys, so a trace of anything else cannot @@ -391,24 +558,23 @@ enum class LookupResult { Miss, Hit, Unsupported }; // original form: attn_scale reads 1, ragged num_tokens read 0, and max_seqlen/batch_size // read their bucketed values. Recover those from the caller if a line needs to be traced // back to a specific test case. -inline void record_cache_lookup(const char *pass, LookupResult result, const FusedAttnConfig &key) { +inline void record_cache_lookup(Backend b, Pass p, LookupResult result, + const FusedAttnConfig &key) { if (!enabled()) return; register_summary_once(); - const bool is_fwd = std::strcmp(pass, "fwd") == 0; - // Unsupported counts as a hit: what the hit column measures is lookups that were answered - // without building anything, and a remembered refusal is one of those. Which kind of answer - // it was shows in the trace line, and the running total of refusals is the unsup column. - const bool hit = (result != LookupResult::Miss); - EventCounters &pc = counters(is_fwd); - (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); - EventCounters &tpc = thread_counters(is_fwd); - (hit ? tpc.hit : tpc.miss).fetch_add(1, std::memory_order_relaxed); + // A refusal replayed from the negative cache is counted apart from a graph hit, in + // hit_unsupported rather than in hit_supported. Both were answered without building + // anything, which is what the two hit columns have in common; which map answered is the + // thing worth being able to read off a level-1 summary, since a run whose hits are mostly + // replayed refusals is not reusing graphs at all. + lookup_column(counters(b, p), result).fetch_add(1, std::memory_order_relaxed); + lookup_column(thread_counters(b, p), result).fetch_add(1, std::memory_order_relaxed); // The per-lookup config dump is the highest-volume line (one per cache lookup); // keep it out of the level-1 path and off the stderr lock unless tracing. if (!trace_enabled()) return; std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %s | %-3s %-5s | tid=%u dev=%d | train=%d det=%d cg=%d " + "[FUSED-ATTN-CACHE] %s%-3s %-3s %-11s | tid=%u dev=%d | 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 @@ -418,26 +584,24 @@ inline void record_cache_lookup(const char *pass, LookupResult result, const Fus " 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", - process_tag().c_str(), pass, - result == LookupResult::Miss ? "MISS" : (result == LookupResult::Hit ? "HIT" : "NOSUP"), - thread_seq_id(), key.device_id, static_cast(key.is_training), - static_cast(key.deterministic), static_cast(key.cuda_graph), - static_cast(key.return_max_logit), static_cast(key.check_for_forward_support), - static_cast(key.attn_mask_type), static_cast(key.bias_type), - static_cast(key.window_size_left), static_cast(key.window_size_right), - static_cast(key.bottom_right_diagonal), static_cast(key.softmax_type), - static_cast(key.scaling_mode), static_cast(key.dropout), - static_cast(key.attn_scale), static_cast(key.qkv_dtype), - static_cast(key.o_dtype), static_cast(key.do_dtype), - static_cast(key.dqkv_dtype), static_cast(key.qkv_layout), - static_cast(key.o_format), static_cast(key.do_format), - static_cast(key.dqkv_layout), static_cast(key.qkv_scale_inv_format), - static_cast(key.do_scale_inv_format), static_cast(key.batch_size), - static_cast(key.num_attn_heads), static_cast(key.num_gqa_groups), - static_cast(key.head_dim_qk), static_cast(key.head_dim_v), - static_cast(key.max_seqlen_q), static_cast(key.max_seqlen_kv), - static_cast(key.num_tokens_q), static_cast(key.num_tokens_kv), - static_cast(key.bucketed_batch_size), + rank_tag().c_str(), backend_name(b), pass_name(p), lookup_name(result), thread_seq_id(), + key.device_id, static_cast(key.is_training), static_cast(key.deterministic), + static_cast(key.cuda_graph), static_cast(key.return_max_logit), + static_cast(key.check_for_forward_support), static_cast(key.attn_mask_type), + static_cast(key.bias_type), static_cast(key.window_size_left), + static_cast(key.window_size_right), static_cast(key.bottom_right_diagonal), + static_cast(key.softmax_type), static_cast(key.scaling_mode), + static_cast(key.dropout), static_cast(key.attn_scale), + static_cast(key.qkv_dtype), static_cast(key.o_dtype), + static_cast(key.do_dtype), static_cast(key.dqkv_dtype), + static_cast(key.qkv_layout), static_cast(key.o_format), + static_cast(key.do_format), static_cast(key.dqkv_layout), + static_cast(key.qkv_scale_inv_format), static_cast(key.do_scale_inv_format), + static_cast(key.batch_size), static_cast(key.num_attn_heads), + static_cast(key.num_gqa_groups), static_cast(key.head_dim_qk), + static_cast(key.head_dim_v), static_cast(key.max_seqlen_q), + static_cast(key.max_seqlen_kv), static_cast(key.num_tokens_q), + static_cast(key.num_tokens_kv), static_cast(key.bucketed_batch_size), static_cast(key.bucketed_num_tokens_q), static_cast(key.bucketed_num_tokens_kv), static_cast(key.num_pages_k), static_cast(key.num_pages_v), static_cast(key.page_size_k), @@ -482,15 +646,16 @@ struct StageTiming { std::atomic time_ns{0}; }; -// Bucketed by pass, so the summary can report the cost of each stage separately -// for forward and backward. Unlike the thread registry above, this table needs no -// leak to outlive the exit handler that reads it: it holds nothing but atomics, so -// it is trivially destructible and no destructor is registered for it at all. -constexpr size_t kStageBuckets = 2 * static_cast(BuildStage::kCount); -inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { +// Bucketed by build site, so the summary can report the cost of each stage separately for each +// backend and pass -- an fp8 build and an f16 build are different work, and averaging them +// together would describe neither. Unlike the thread registry above, this table needs no leak to +// outlive the exit handler that reads it: it holds nothing but atomics, so it is trivially +// destructible and no destructor is registered for it at all. +constexpr size_t kStageBuckets = kSiteCount * static_cast(BuildStage::kCount); +inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { static std::array table{}; const size_t idx = - (is_fwd ? 0u : 1u) * static_cast(BuildStage::kCount) + static_cast(s); + site_index(b, p) * static_cast(BuildStage::kCount) + static_cast(s); return table[idx]; } @@ -506,9 +671,10 @@ inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { struct ScopedBuildTimer { BuildStage stage; bool on; - bool is_fwd; + Backend backend; + Pass pass; std::chrono::steady_clock::time_point start; - ScopedBuildTimer(bool is_fwd_, BuildStage s) : stage(s), on(enabled()), is_fwd(is_fwd_) { + ScopedBuildTimer(Backend b, Pass p, BuildStage s) : stage(s), on(enabled()), backend(b), pass(p) { if (!on) return; register_summary_once(); start = std::chrono::steady_clock::now(); @@ -519,20 +685,19 @@ struct ScopedBuildTimer { static_cast(std::chrono::duration_cast( std::chrono::steady_clock::now() - start) .count()); - StageTiming &t = stage_timing(is_fwd, stage); + StageTiming &t = stage_timing(backend, pass, stage); t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); t.calls.fetch_add(1, std::memory_order_relaxed); } }; -// Time `fn` as `stage` of the given pass ("fwd"/"bwd", matching the record_* -// helpers above). Preferred over declaring a ScopedBuildTimer at the call site: -// the measured region is exactly the call passed in, so surrounding work cannot -// drift into it as that code changes. With diagnostics off this costs the pass -// comparison and one cached-flag check; both are per build, not per lookup. +// Time `fn` as `stage` of the given build site, named as the record_* helpers above name it. +// Preferred over declaring a ScopedBuildTimer at the call site: the measured region is exactly +// the call passed in, so surrounding work cannot drift into it as that code changes. With +// diagnostics off this costs one cached-flag check, and that is per build rather than per lookup. template -inline void timer(const char *pass, BuildStage stage, Fn &&fn) { - ScopedBuildTimer scoped(std::strcmp(pass, "fwd") == 0, stage); +inline void timer(Backend b, Pass p, BuildStage stage, Fn &&fn) { + ScopedBuildTimer scoped(b, p, stage); fn(); } @@ -547,9 +712,19 @@ inline void register_summary_once() { // that the blocks of concurrently-exiting processes (one per rank under // torchrun) stay grouped instead of interleaving line by line. std::string block; - block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary begin =====\n"; - // Per-thread breakdown (sorted by tid). Useful in the single-process - // context-parallel case where each device runs on its own thread. + block += "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary begin =====\n"; + constexpr Backend kBackends[] = {Backend::F16, Backend::FP8}; + // A backend the run never reached is left out of the summary rather than reported as a + // row of zeros, so the usual single-backend run reads as it did before this was split. + size_t active_backends = 0; + for (const Backend b : kBackends) { + if (!snapshot(counters(b, Pass::Fwd)).empty() || + !snapshot(counters(b, Pass::Bwd)).empty()) { + ++active_backends; + } + } + // Per-thread breakdown (sorted by tid), one row per backend that thread drove. Useful in + // the single-process context-parallel case where each device runs on its own thread. { std::lock_guard lock(thread_registry_mutex()); std::vector blocks = thread_registry(); @@ -561,31 +736,53 @@ inline void register_summary_once() { std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); std::snprintf(dev_field, sizeof(dev_field), "dev=%d", tc->device.load(std::memory_order_relaxed)); - block += format_counter_line("SUMMARY-TID", tid_field, dev_field, tc->fwd, tc->bwd); + for (const Backend b : kBackends) { + const CounterSnapshot fwd = snapshot(tc->sites[site_index(b, Pass::Fwd)]); + const CounterSnapshot bwd = snapshot(tc->sites[site_index(b, Pass::Bwd)]); + if (fwd.empty() && bwd.empty()) continue; + char label[32]; + std::snprintf(label, sizeof(label), "%s SUMMARY-TID", backend_name(b)); + block += format_counter_line(label, tid_field, dev_field, fwd, bwd); + } } } - // Totals last, so they read as the sum of the per-thread lines above. - block += format_counter_line("SUMMARY", "tid=all", "dev=all", counters(/*is_fwd=*/true), - counters(/*is_fwd=*/false)); - 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; - char line[288]; - std::snprintf(line, sizeof(line), - "[FUSED-ATTN-CACHE] %s | %-3s %-22s | calls=%" PRIu64 - " | time=%9.3f ms/call\n", - process_tag().c_str(), pass, kStageNames[i], n, total_ms / n); - block += line; + // Totals last, so they read as the sum of the per-thread rows above: one row per backend, + // then a row across the backends only when the run used more than one. With a single + // backend that row would repeat the one above it verbatim and say nothing extra. + CounterSnapshot all_fwd; + CounterSnapshot all_bwd; + for (const Backend b : kBackends) { + const CounterSnapshot fwd = snapshot(counters(b, Pass::Fwd)); + const CounterSnapshot bwd = snapshot(counters(b, Pass::Bwd)); + all_fwd += fwd; + all_bwd += bwd; + if (fwd.empty() && bwd.empty()) continue; + char label[32]; + std::snprintf(label, sizeof(label), "%s SUMMARY", backend_name(b)); + block += format_counter_line(label, "tid=all", "dev=all", fwd, bwd); + } + if (active_backends > 1) { + block += format_counter_line("SUMMARY", "tid=all", "dev=all", all_fwd, all_bwd); + } + for (const Backend b : kBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + const BuildStage s = static_cast(i); + const StageTiming &t = stage_timing(b, p, 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; + char line[288]; + std::snprintf( + line, sizeof(line), + "[FUSED-ATTN-CACHE] %s%-3s %-3s %-22s | calls=%" PRIu64 " | time=%9.3f ms/call\n", + rank_tag().c_str(), backend_name(b), pass_name(p), kStageNames[i], n, total_ms / n); + block += line; + } } } - block += "[FUSED-ATTN-CACHE] " + process_tag() + " | ===== summary end =====\n"; + block += "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary end =====\n"; std::fwrite(block.data(), 1, block.size(), stderr); std::fflush(stderr); }); diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 391a22768a..d1bbeae4ad 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -223,6 +223,11 @@ void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int6 // by the cu_seqlens_padded_to_offsets conversion kernel and the direct-seqlens path // (which passes them to cuDNN as ragged offset multipliers). struct RaggedOffsetMultipliers { + // Zeroed, for a FusedAttnConfig that has not been through derive() yet. Every multiplier is + // a per-token element count, so zero is not a usable value; it is only ever read after + // derive() has replaced it, which check_derived() is what enforces. + RaggedOffsetMultipliers() = default; + RaggedOffsetMultipliers(NVTE_QKV_Layout_Group layout_group, int64_t h, int64_t hg, int64_t d_qk, int64_t d_v) : q(h * d_qk), k(hg * d_qk), v(hg * d_v), o(h * d_v), stats(h), kv_from_q(false) { @@ -241,13 +246,13 @@ struct RaggedOffsetMultipliers { } } - int64_t q; - int64_t k; - int64_t v; - int64_t o; - int64_t stats; + int64_t q = 0; + int64_t k = 0; + int64_t v = 0; + int64_t o = 0; + int64_t stats = 0; // K/V offsets scale the Q-side cu_seqlens_padded (interleaved QKV layouts) - bool kv_from_q; + bool kv_from_q = false; }; __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index d2307b07b6..d7106803d0 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -414,12 +414,12 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, /*! \brief Get fused-attention backend based on user configuration. * * 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``). + * and returns a backend if supported, otherwise a message explaining why the configuration is not supported. + * If supported, the backend is cached and reused for future calls with the same configuration. * * \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; + * \param[out] message If the configuration is supported, an empty string; * 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 @@ -1002,10 +1002,6 @@ class FusedAttnConfigWrapper { FusedAttnConfigWrapper &operator=(FusedAttnConfigWrapper &&other) noexcept { if (this != &other) { - // Guarded as the destructor is. A moved-from wrapper holds nullptr, and the C API rejects a - // NULL handle by throwing; thrown out of a noexcept function that is a call to - // std::terminate, which no caller can catch. The guard belongs on this side rather than in - // nvte_destroy_*, so that the C entry point keeps reporting a genuinely bad handle. if (cfg_ != nullptr) { nvte_destroy_fused_attn_config(cfg_); } @@ -1183,8 +1179,6 @@ class FusedAttnFwdParamsWrapper { FusedAttnFwdParamsWrapper &operator=(FusedAttnFwdParamsWrapper &&other) noexcept { if (this != &other) { - // See FusedAttnConfigWrapper::operator=: destroying a moved-from (NULL) handle throws out - // of a noexcept function, which is std::terminate. if (params_ != nullptr) { nvte_destroy_fused_attn_fwd_params(params_); } @@ -1335,8 +1329,6 @@ class FusedAttnBwdParamsWrapper { FusedAttnBwdParamsWrapper &operator=(FusedAttnBwdParamsWrapper &&other) noexcept { if (this != &other) { - // See FusedAttnConfigWrapper::operator=: destroying a moved-from (NULL) handle throws out - // of a noexcept function, which is std::terminate. if (params_ != nullptr) { nvte_destroy_fused_attn_bwd_params(params_); } From 63b959ac26c78f2ee4b532c25738537c7926d635 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 19 Aug 2026 09:06:02 -0700 Subject: [PATCH 82/88] WIP: clean up graph cache/debug Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 8 +- docs/examples/attention/attention.ipynb | 8 +- tests/pytorch/attention/test_attention.py | 39 +- .../common/fused_attn/config_and_params.cpp | 72 ++- .../common/fused_attn/config_and_params.h | 59 +- .../common/fused_attn/fused_attn.cpp | 10 +- .../fused_attn_f16_arbitrary_seqlen.cu | 213 ++----- .../common/fused_attn/fused_attn_fp8.cu | 292 ++++----- .../common/fused_attn/graph_cache.h | 462 +++++--------- .../common/fused_attn/graph_cache_debug.h | 587 +++++++----------- 10 files changed, 716 insertions(+), 1034 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index af06ebbb58..f17842c5e1 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -202,13 +202,13 @@ backend-selection overview. :Default: ``0`` :Description: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). Output goes to stderr, prefixed ``[FUSED-ATTN-CACHE]``. - ``1`` emits one line per event that happens once per distinct cache key -- ``BUILD_GRAPH`` when a graph is constructed, ``BUILD_PLANS`` when its kernels are compiled on first execution, ``UNSUPPORTED`` when cuDNN refuses a configuration -- plus an end-of-run ``SUMMARY`` (per backend, per thread, and across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries. + ``1`` emits one line per event that happens once per distinct cache key -- ``CREATE_GRAPH`` when a graph is constructed, ``BUILD_PLANS`` when its kernels are compiled on first execution -- plus an end-of-run summary block (one row per build site, per thread and in total, plus a row per pass across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries. - Every line names the build site behind it -- ``f16`` or ``fp8``, then the pass -- and carries only that backend's counters, so a process that uses both can still tell which of them built what. A backend the run never reached is left out of the summary entirely. + Every line names the thread and device it came from, then the build site behind it -- ``f16`` or ``fp8``, then the pass -- and carries the counters of that site alone, so a process that uses both backends can still tell which of them built what. One line is one pass; the forward and the backward read as adjacent rows. A summary row is an event line without the event name, the block is delimited by ``===== summary begin =====`` and ``===== summary end =====``, and ``tid=all dev=all`` marks the totals. A build site the run never reached is left out rather than shown as a row of zeros. - The two hit columns name which of the cache's two maps answered a lookup: ``hit_supported`` is a cached graph reused, while ``hit_unsupported`` is a configuration cuDNN already refused, replayed from the negative cache rather than rebuilt. A run whose hits are mostly ``hit_unsupported`` is not reusing graphs at all -- it is asking repeatedly for something that will never run fused. Together with ``miss`` these account for every lookup, and ``unsupported`` counts the refusals themselves, so it stays at one per bad configuration however many times that configuration is queried. + ``hit`` and ``miss`` account for every lookup, and a miss ends either in ``create_graph`` or in a build cuDNN refused, so ``miss`` minus ``create_graph`` is the number of refusals. Nothing is cached for a configuration cuDNN refuses, so that difference counts refused builds rather than refused configurations: a configuration that is queried again is built and refused again. ``miss`` climbing while ``create_graph`` stays put says this site never runs fused and keeps paying a discarded graph build to find that out, which makes it the pair to read when attention is slower than expected and nothing raised an error. The reason cuDNN gave is not logged here; it reaches the framework as the message explaining why the fused backend was not selected. - ``2`` additionally emits a per-lookup ``HIT``/``MISS``/``UNSUPPORTED`` line carrying the full cache key, and a per-execution ``EXEC`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. A level-2 ``UNSUPPORTED`` is a lookup answered from a stored refusal, as opposed to the level-1 event that recorded it; the counter line carries totals, the lookup line carries the key. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. + ``2`` additionally emits a per-lookup ``HIT``/``MISS`` line carrying the full cache key, and a per-execution ``EXEC`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 1c206264f9..7f5deee722 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -259,13 +259,13 @@ "```\n", "NVTE_FUSED_ATTN_CACHE_DEBUG = 0/1/2 # disables/enables graph cache diagnostics\n", "```\n", - "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]` and names the build site behind it -- `f16` or `fp8`, then the pass -- and there is one per event that happens once per configuration: `BUILD_GRAPH` when a graph is constructed, `BUILD_PLANS` when its kernels are compiled on first execution, and `UNSUPPORTED` when cuDNN declines a configuration. Each event name is also the counter column it increments, and the two hit columns say which of the cache's maps answered a lookup: `hit_supported` is a cached graph reused, `hit_unsupported` a configuration cuDNN had already refused. An end-of-run `SUMMARY` gives the totals per backend, thread and device, followed by where the build time went:\n", + "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]`, names the thread and device it came from and then the build site behind it -- `f16` or `fp8`, then the pass -- and there is one per event that happens once per configuration: `CREATE_GRAPH` when a graph is constructed, and `BUILD_PLANS` when its kernels are compiled on first execution. Each event name is also the counter column it increments. A configuration cuDNN declines raises a miss and no build, so `miss` minus `create_graph` counts refused builds; nothing is cached for a refusal, so a configuration that is queried again is built and refused again. An end-of-run summary block gives one row per build site, per thread and in total (`tid=all dev=all`), followed by where the build time went:\n", "```\n", - "[FUSED-ATTN-CACHE] f16 fwd BUILD_GRAPH | tid=0 dev=0 | fwd hit_supported=0, miss=1, build_graph=1, ...\n", - "[FUSED-ATTN-CACHE] f16 SUMMARY | tid=all dev=all | fwd hit_supported=5, miss=1, build_graph=1, ...\n", + "[FUSED-ATTN-CACHE] tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, ...\n", + "[FUSED-ATTN-CACHE] tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, ...\n", "[FUSED-ATTN-CACHE] f16 fwd check_support | calls=1 | time= 42.135 ms/call\n", "```\n", - "The number to read first is `build_graph`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n", + "The number to read first is `create_graph`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n", "\n", "This flag is supported in both PyTorch and JAX, since the cache it reports on lives in the common C++ layer. By default only rank 0 emits; see [NVTE_FUSED_ATTN_CACHE_DEBUG](../../envvars.rst) for selecting other ranks." ] diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index b5413d6687..42bb3f0e9f 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -275,18 +275,19 @@ def test_dpa_checkpoint(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, None, False, False) -# One [FUSED-ATTN-CACHE] event, as either a counter line ("f16 fwd BUILD_GRAPH") or a level-2 -# trace line ("f16 fwd MISS"). Both name the build site first, and the backend half of it is -# matched but not counted on: the worker below drives whichever one its dtype selects, and every -# assertion here holds of either. The pass and the event name are what this test reads, plus the -# trace line's cache key, kept so that distinct keys can be counted. Every event name is also the -# counter column it increments, except UNSUPPORTED, which the diagnostics use for both the -# level-1 refusal and the level-2 lookup answered from it -- both count as refusals here. The -# rank prefix is optional because it is emitted only when the launcher exports a rank, which a -# plain subprocess like the worker does not. +# One [FUSED-ATTN-CACHE] event, as either a counter line ("f16 fwd CREATE_GRAPH") or a level-2 +# trace line ("f16 fwd MISS"). Both name the thread and device first and then the build site, of +# which the backend half is matched but not counted on: the worker below drives whichever one its +# dtype selects, and every assertion here holds of either. The pass and the event name are what +# this test reads, plus the trace line's cache key, kept so that distinct keys can be counted. +# Every event name is also the counter column it increments. Requiring an event name is also what +# excludes the end-of-run summary rows, which are otherwise the same shape. The rank prefix is +# optional because it is emitted only when the launcher exports a rank, which a plain subprocess +# like the worker does not. _CACHE_EVENT = re.compile( - r"\[FUSED-ATTN-CACHE\]\s+(?:rank=\d+\s+\|\s+)?(?Pf16|fp8)\s+(?Pfwd|bwd)\s+" - r"(?PBUILD_GRAPH|BUILD_PLANS|UNSUPPORTED|EXEC|MISS|HIT)\b(?P.*)" + r"\[FUSED-ATTN-CACHE\]\s+(?:rank=\d+\s+\|\s+)?tid=\d+\s+dev=-?\d+\s+\|\s+" + r"(?Pf16|fp8)\s+(?Pfwd|bwd)\s+" + r"(?PCREATE_GRAPH|BUILD_PLANS|EXEC|MISS|HIT)\b(?P.*)" ) _CACHE_PHASE = re.compile(r"\[CACHE-TEST\] phase=(?P\w+)") @@ -371,15 +372,17 @@ def count(phase, event, pass_name=pass_name): # The first query builds each pass's graph, and no more than its graph: a support # query stops at check_support(), leaving the kernel compilation (BUILD_PLANS) to - # whoever executes it. + # whoever executes it. A build cuDNN refused would show up here as the miss without + # the build, since nothing is recorded for a refusal. assert count("query", "MISS") == 1, f"{pass_name}: expected one cold miss{context}" - assert count("query", "BUILD_GRAPH") == 1, f"{pass_name}: expected one build{context}" - assert count("query", "UNSUPPORTED") == 0, f"{pass_name}: cuDNN refused the config{context}" + assert count("query", "CREATE_GRAPH") == 1, f"{pass_name}: expected one build{context}" assert count("query", "BUILD_PLANS") == 0, f"{pass_name}: query compiled kernels{context}" # Asking the identical question again must cost nothing. assert count("requery", "MISS") == 0, f"{pass_name}: repeated query missed{context}" - assert count("requery", "BUILD_GRAPH") == 0, f"{pass_name}: repeated query rebuilt{context}" + assert ( + count("requery", "CREATE_GRAPH") == 0 + ), f"{pass_name}: repeated query rebuilt{context}" assert count("requery", "HIT") >= 1, f"{pass_name}: repeated query never looked{context}" # The execution must find the graph the query left behind -- a miss here is the @@ -389,7 +392,7 @@ def count(phase, event, pass_name=pass_name): count("exec", "MISS") == 0 ), f"{pass_name}: execution missed the query's graph{context}" assert ( - count("exec", "BUILD_GRAPH") == 0 + count("exec", "CREATE_GRAPH") == 0 ), f"{pass_name}: execution rebuilt the graph{context}" assert count("exec", "EXEC") >= 1, f"{pass_name}: fused attention never ran{context}" assert count("exec", "BUILD_PLANS") == 1, f"{pass_name}: expected one plan build{context}" @@ -398,7 +401,7 @@ def count(phase, event, pass_name=pass_name): # a different scale has to reuse everything, down to the compiled kernels. assert count("rescale", "MISS") == 0, f"{pass_name}: attn_scale changed the key{context}" assert ( - count("rescale", "BUILD_GRAPH") == 0 + count("rescale", "CREATE_GRAPH") == 0 ), f"{pass_name}: attn_scale forced a build{context}" assert count("rescale", "BUILD_PLANS") == 0, f"{pass_name}: attn_scale recompiled{context}" assert count("rescale", "EXEC") >= 1, f"{pass_name}: rescaled run did not execute{context}" @@ -406,7 +409,7 @@ def count(phase, event, pass_name=pass_name): # max_seqlen is a dimension the graph is built at, so it must miss -- once, for one # new graph, rather than invalidating what is already cached. assert count("reshape", "MISS") == 1, f"{pass_name}: expected one miss{context}" - assert count("reshape", "BUILD_GRAPH") == 1, f"{pass_name}: expected one build{context}" + assert count("reshape", "CREATE_GRAPH") == 1, f"{pass_name}: expected one build{context}" assert ( len(miss_keys["reshape"][pass_name]) == 1 ), f"{pass_name}: more than one new cache key{context}" diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 451280b2c9..e6ee965b02 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -125,11 +125,41 @@ void FusedAttnConfig::derive() { is_derived = true; } -FusedAttnConfig FusedAttnConfig::make_cache_key() const { +GraphDims graph_dims(const FusedAttnConfig &cfg, Pass pass) { + check_derived(cfg); + GraphDims dims; + + // The one condition both answers turn on: the forward graph can be handed the user's cu_seqlens* + // buffers untouched, and then it is those buffers the graph has to match -- their + // [batch_size + 1] length, which a bucketed batch would read past the end of, and their int32 + // width. The backward graph always reads seqlens converted into our own workspace, so nothing + // there is sized by the true batch and nothing there is held to int32. + const bool cudnn_reads_users_cu_seqlens = pass == Pass::Fwd && cfg.uses_cu_seqlens_directly; + + if (cudnn_reads_users_cu_seqlens) { + dims.ragged_offset_type = DType::kInt32; + } else { + // Choose between 32-bit and 64-bit offsets by what the runtime supports, which is what lets + // older cuDNN runtimes work rather than fail. + dims.ragged_offset_type = cudnnGetVersion() >= 90500 ? DType::kInt64 : DType::kInt32; + } + + // Build at the bucketed batch where a ragged layout is packed, so that one graph serves every + // batch in its bucket -- the same reason graph_max_seqlen_* stands in for the sequence lengths. + dims.batch_size = static_cast(cfg.batch_size); + if ((cfg.is_ragged_q || cfg.is_ragged_kv) && cfg.uses_packed_ragged_graph && + !cudnn_reads_users_cu_seqlens) { + dims.batch_size = static_cast(cfg.bucketed_batch_size); + } + + return dims; +} + +FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { // Requires a derived config: every normalization below reads a derived field -- is_padding and // is_causal_bottom_right, the is_ragged_* pair, the graph_max_seqlen_* dimensions, and the - // uses_* flags. A precondition rather than an assert, since all four callers construct their - // GraphInputs first and that constructor asserts it. + // uses_* flags. A precondition rather than an assert, since every caller reaches this through a + // cache_graph_* wrapper, which asserts it once for both the key and the graph. FusedAttnConfig cache_cfg = *this; // Key the device ID for multi-GPU single-process runs @@ -152,18 +182,14 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { cache_cfg.max_seqlen_q = cache_cfg.graph_max_seqlen_q; cache_cfg.max_seqlen_kv = cache_cfg.graph_max_seqlen_kv; - // Bucket the THD (ragged) batch, and drop the token counts the bucketing has replaced + // Name the batch size the graph is built at, and drop the token counts the bucketing replaced. + // Asking graph_dims() rather than restating its rule is what keeps the key from naming a batch + // the graph was not built with -- the two directions bucket differently, and the graph builders + // ask the same question with the same pass. if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_packed_ragged_graph) { cache_cfg.num_tokens_q = 0; cache_cfg.num_tokens_kv = 0; - // The forward graph keeps the true batch size when it takes the user's cu_seqlens - // directly, since cuDNN reads those [actual_b+1] buffers itself; the backward graph - // converts them and so always buckets. The key has to follow whichever the graph does, - // or it would name a batch size the graph was not built with. See F16BwdGraphInputs. - const bool bucket_batch = !check_for_forward_support || !cache_cfg.uses_cu_seqlens_directly; - if (bucket_batch) { - cache_cfg.batch_size = cache_cfg.bucketed_batch_size; - } + cache_cfg.batch_size = static_cast(graph_dims(*this, pass).batch_size); } // attn_scale is a pass-by-value graph input and different scales can share the same cached graph @@ -175,20 +201,28 @@ FusedAttnConfig FusedAttnConfig::make_cache_key() const { // give a workload that both captures and runs eagerly two entries for every configuration. cache_cfg.cuda_graph = false; - // Restrict each direction's key to the fields its graph actually consumes, so - // no redundant graphs are built and no cache misses either - if (check_for_forward_support && !check_for_backward_support) { + // Restrict this direction's key to the fields its graph actually consumes, so no redundant + // graphs are built and no cache misses either. Keyed on the pass rather than on the + // check_for_*_support flags, so that a caller asking about both directions -- which every + // backend query from a framework does -- still gets a key each pass can find its own graph + // under, instead of one narrowed for neither. + if (pass == Pass::Fwd) { 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; - } - if (check_for_backward_support && !check_for_forward_support) { + } else { cache_cfg.return_max_logit = false; } + // The two flags say which directions the caller wanted probed, which the graph this key names + // does not depend on. Normalized so that a key is the same whether it came from a probe or from + // execution, and so that a level-2 trace line cannot claim a direction the key is not for. + cache_cfg.check_for_forward_support = pass == Pass::Fwd; + cache_cfg.check_for_backward_support = pass == Pass::Bwd; + return cache_cfg; } @@ -303,8 +337,8 @@ FusedAttnConfig FusedAttnFwdParams::make_config() const { FusedAttnConfig FusedAttnBwdParams::make_config() const { const FusedAttnBwdParams ¶ms = *this; FusedAttnConfig cfg{}; - // Backward execution: only the backward graph is run. check_for_forward_support=false also - // selects the backward key normalization in make_cache_key(). + // Backward execution: only the backward graph is run, so do not pay for a forward support + // check whose graph this call will never execute. cfg.check_for_forward_support = false; cfg.check_for_backward_support = true; cfg.is_training = true; diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 44e0ef457b..85f582bd19 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -20,6 +20,16 @@ namespace transformer_engine { namespace fused_attn { +// Which of the two graphs a config is being turned into. Declared here, rather than with the +// diagnostics that also name it, because two things about a config depend on the direction: +// make_cache_key() below, and graph_dims() further down. +// +// Passed in rather than derived, because a config cannot say which graph is being built from it. +// check_for_forward_support and check_for_backward_support state which directions a caller wants +// probed, and a backend query arriving from a framework has both set, so they answer a different +// question -- see the comment on them below. +enum class Pass { Fwd, Bwd }; + struct FusedAttnConfig { // basic attention settings bool is_training = true; @@ -82,6 +92,11 @@ 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. check_for_forward_support). Added for convinence // purposes and do not represent any graph properties. + // + // The two below say which directions nvte_get_fused_attn_backend_v2 should probe, and nothing + // else: not which graph is being built, which is what Pass names. They default to true and the + // attribute API cannot reach them, so a backend query from a framework asks about both + // directions, while the execution entry points set the one they are about to run. bool check_for_forward_support = true; bool check_for_backward_support = true; // Whether derive() has run, i.e. whether the fields below hold anything. Every consumer of a @@ -91,7 +106,7 @@ struct FusedAttnConfig { // derive() recomputes unconditionally, so a config whose inputs change can simply be re-derived. bool is_derived = false; // THD batch/token counts, the raw buckets. The graph dimensions built out of them are - // graph_max_seqlen_* below and, because the batch is direction-dependent, F16FwdGraphInputs::b. + // graph_max_seqlen_* below and, because the batch is direction-dependent, graph_dims(). size_t bucketed_batch_size = 0; size_t bucketed_num_tokens_q = 0; size_t bucketed_num_tokens_kv = 0; @@ -99,9 +114,9 @@ struct FusedAttnConfig { bool uses_cu_seqlens_directly = false; // Whether a ragged (THD) graph is built at packed token-count dimensions with ragged Stats/LSE, // rather than at dense max_seqlen ones. Held here rather than asked for at each of the places - // that need it -- graph_max_seqlen_* below, make_cache_key()'s batch, and the two GraphInputs -- - // because the key and the graph have to be built at the same dimensions, and two independent - // queries are two chances to disagree. Unlike the flags above, this one depends on the device as + // that need it -- graph_max_seqlen_* below, and graph_dims() -- because the key and the + // graph have to be built at the same dimensions, and two independent queries are two chances to + // disagree. Unlike the flags above, this one depends on the device as // well as the cuDNN version, so a config carries the answer for the device it was derived on; // every entry point derives immediately before use, and the cache key records device_id. bool uses_packed_ragged_graph = false; @@ -116,8 +131,8 @@ struct FusedAttnConfig { // dimensions the graph was built with -- a key that says otherwise is a hit on a graph of the // wrong shape -- and stating the substitution once is what keeps make_cache_key() and the graph // builders from drifting. Both passes build at the same sequence lengths; the batch size is the - // one dimension they disagree on, so it stays with the direction that knows, in - // F16FwdGraphInputs and F16BwdGraphInputs. + // one dimension they disagree on, which is why that one comes from graph_dims() below instead of + // a field here. size_t graph_max_seqlen_q = 0; size_t graph_max_seqlen_kv = 0; // Elements per token for each ragged tensor, from the layout group and the head dimensions. @@ -240,7 +255,11 @@ struct FusedAttnConfig { // 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. - FusedAttnConfig make_cache_key() const; + // + // `pass` is which graph the key is for. It decides both direction-dependent normalizations -- + // which fields are dropped, and whether the batch is bucketed -- so a key built for one pass + // cannot be handed to the other's cache. + FusedAttnConfig make_cache_key(Pass pass) const; }; // Assert that `cfg` has been through derive(), for code about to read a derived field. Worth @@ -258,6 +277,32 @@ inline void check_derived(const FusedAttnConfig &cfg) { "fused_attn.cpp."); } +// What a graph is built with that depends on which direction it is for, and so cannot be fields +// derive() fills: one stored value cannot answer for both passes, and the selector derives a config +// once and probes both directions off it. +struct GraphDims { + // The batch size the graph is built at: the bucketed one where a ragged layout is packed, so + // that one graph serves every batch in its bucket, and the true one otherwise. + int64_t batch_size = 0; + // The width the graph expects ragged (THD) offsets in. + DType ragged_offset_type = DType::kInt32; +}; + +// Both of the above for one direction. One function returning the pair rather than two returning +// one each, because a single condition decides both -- see the body -- and stating that condition +// twice is what would let the two drift into a graph built at a bucketed batch that reads its +// offsets at the other width. +// +// A free function rather than a member for the same reason as check_derived() above: it reads +// public fields and computes, and is no part of the config's own invariants. Requires a derived +// config and asserts it. +// +// Everything that has to agree about a graph asks this with the same pass: the builder, the code +// that binds runtime pointers to the built graph, and make_cache_key(). That is the point of +// having one place to ask, since a disagreement means a graph built at dimensions the pointers +// bound to it do not describe. +GraphDims graph_dims(const FusedAttnConfig &cfg, Pass pass); + 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/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index a7c5e25ef9..bd016e013e 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -460,14 +460,14 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // | | `-- reject -> NVTE_No_Backend + reason -> the NVTE_ERROR below // | | // | `-- is_supported_f16_fwd / is_supported_fp8_fwd -// | `-- f16_fwd_cached_graph(): builds and inserts the entry, or throws -// | UnsupportedGraph, whose message becomes the reason for the refusal +// | `-- cache_graph_f16_fwd(): builds and inserts the entry, or throws, in +// | which case cuDNN's message becomes the reason for the refusal // | // `-- fused_attn_arbitrary_seqlen_fwd -> ..._fwd_impl the selected backend // | -// +-- f16_fwd_cached_graph() HIT: the entry the query above just built -// +-- ensure_plans_built() the kernel compilation, once per entry -// `-- bind device pointers, execute() +// +-- cache_graph_f16_fwd() HIT: the entry the query above just built +// +-- build_plans() the kernel compilation, once per entry +// `-- bind device pointers, graph.execute() void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; 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 42130de8ff..63094fa8c4 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 @@ -26,9 +26,9 @@ namespace fused_attn { namespace fe = cudnn_frontend; // Every graph-cache event raised here names the build site it came from. This file is the f16 -// arbitrary-seqlen backend throughout; only the pass differs between call sites. +// arbitrary-seqlen backend throughout; only the pass differs between call sites. Pass itself needs +// no using-declaration: it is fused_attn::Pass, since the config answers by direction too. using graph_cache_debug::Backend; -using graph_cache_debug::Pass; using SdpaF16FwdGraphAndTensors = std::tuple, @@ -53,62 +53,19 @@ using SdpaF16FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// What the forward graph is built from that the config cannot say on its own, because the answer -// depends on the direction: the batch size, and the width the ragged offsets are written in. The -// sequence lengths are not here -- both passes build at cfg.graph_max_seqlen_* -- and neither is -// anything else a backward graph would answer the same way. The build and the execution have to -// agree on all of it, otherwise the graph is built for different dimensions than the pointers -// bound to it describe, or with a ragged offset width the offsets are not written in, so it is -// derived once, here, and handed to both. -struct F16FwdGraphInputs { - // Everything below is arithmetic on an already-derived cfg. Configurations no graph can serve - // are rejected before this point: FusedAttnConfig::derive() asserts them, and - // nvte_get_fused_attn_backend_v2 states them as rules so a support query can answer for them. - explicit F16FwdGraphInputs(const FusedAttnConfig &cfg); - - // The batch size the graph is built at: bucketed for a packed ragged layout, so that one graph - // serves every batch in the same bucket, except when cu_seqlens go to cuDNN directly. - int64_t b = 0; - // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever - // the bucketing above did to `b`. - int64_t actual_b = 0; - DType ragged_offset_type = DType::kInt32; -}; - -F16FwdGraphInputs::F16FwdGraphInputs(const FusedAttnConfig &cfg) { - check_derived(cfg); - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; - const auto cudnn_runtime_version = cudnnGetVersion(); - - b = static_cast(cfg.batch_size); - // keep original batch size because cu_seqlens are created with [b+1] shape - actual_b = b; - // Replace the batch size with the bucketed one so the graph is static within its bucket, the - // same reason cfg.graph_max_seqlen_* replaces the sequence lengths. When passing cu_seqlens* - // directly to cuDNN SDPA, keep the true batch size: cuDNN reads the user's [actual_b+1] - // cu_seqlens buffers, so a quantized batch would read out of bounds. - if ((is_ragged_q || is_ragged_kv) && cfg.uses_packed_ragged_graph && !use_cu_seqlens_directly) { - b = static_cast(cfg.bucketed_batch_size); - } - - ragged_offset_type = - use_cu_seqlens_directly - ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. - : (cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32); -} - // Constructs the forward graph for one cache key, and only constructs it: whether cuDNN will run -// it is settled by the caller, in build_or_get_cached_graph(), which is also where the plan build +// it is settled by the caller, in lookup_or_cache_graph(), which is also where the plan build // eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every call // that does need one now sits on the other side of that boundary. // -// Everything the graph's shape and topology depends on comes from `cfg` and `in`, so the build -// has one source of truth and cannot drift from the caller that will bind pointers to it. -static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig &cfg, - const F16FwdGraphInputs &in) { - const int64_t b = in.b; +// Everything the graph's shape and topology depends on comes from `cfg`, so the build has one +// source of truth and cannot drift from the caller that will bind pointers to it. The two +// dimensions the config cannot answer on its own -- the batch size, and the width ragged offsets +// are written in, both of which differ between the passes -- come from graph_dims() asked with +// Pass::Fwd, the same way the code binding pointers to this graph asks. +static SdpaF16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg) { + const GraphDims dims = graph_dims(cfg, Pass::Fwd); + const int64_t b = dims.batch_size; const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); const cudnn_frontend::DataType_t tensorType = @@ -151,7 +108,7 @@ static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; const auto cudnn_runtime_version = cudnnGetVersion(); const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = in.ragged_offset_type; + const DType ragged_offset_type = dims.ragged_offset_type; const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; const bool generate_stats = true; // Always return stats @@ -420,11 +377,13 @@ static SdpaF16FwdGraphAndTensors build_sdpa_f16_fwd_graph(const FusedAttnConfig // probe come through here, so a probe leaves behind exactly the entry a later execution finds. // That is what lets the probe's answer describe the graph that actually runs, rather than a // separately built lookalike. -static std::shared_ptr> f16_fwd_cached_graph( - const FusedAttnConfig &cfg, const F16FwdGraphInputs &in, cudnnHandle_t handle) { +static std::shared_ptr> cache_graph_f16_fwd( + const FusedAttnConfig &cfg, cudnnHandle_t handle) { static GraphCache cache; - return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::F16, Pass::Fwd, handle, - [&] { return build_sdpa_f16_fwd_graph(cfg, in); }); + // Asserted once here for both the key and the graph, which read the same derived fields. + check_derived(cfg); + return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Fwd), Backend::F16, Pass::Fwd, + handle, [&] { return create_graph_f16_fwd(cfg); }); } void fused_attn_arbitrary_seqlen_fwd_impl( @@ -436,12 +395,14 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - // Derived once and handed to the cache, which passes them to the graph build, so that the - // graph and the pointers bound to it below cannot be decided differently. - const F16FwdGraphInputs in(cfg); - const int64_t b = in.b; - const int64_t actual_b = in.actual_b; - const DType ragged_offset_type = in.ragged_offset_type; + // Asked with the same pass the graph was built with, so that the dimensions below and the ones + // the graph was built at cannot be decided differently. + const GraphDims dims = graph_dims(cfg, Pass::Fwd); + const int64_t b = dims.batch_size; + const DType ragged_offset_type = dims.ragged_offset_type; + // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever the + // bucketing above did to `b`. + const int64_t actual_b = static_cast(cfg.batch_size); const bool use_ragged_stats = cfg.uses_ragged_stats; const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; @@ -461,13 +422,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; try { - auto cache_entry = f16_fwd_cached_graph(cfg, in, handle); + auto cache_entry = cache_graph_f16_fwd(cfg, handle); 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] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built(Backend::F16, Pass::Fwd, *cache_entry); + build_plans(Backend::F16, Pass::Fwd, *cache_entry); // 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. @@ -637,45 +598,11 @@ using SdpaF16BwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The backward equivalent of F16FwdGraphInputs; see there for why these two dimensions are the -// only ones that cannot live on the config, and for why they are derived once and shared. -struct F16BwdGraphInputs { - explicit F16BwdGraphInputs(const FusedAttnConfig &cfg); - - int64_t b = 0; - int64_t actual_b = 0; - DType ragged_offset_type = DType::kInt32; -}; - -F16BwdGraphInputs::F16BwdGraphInputs(const FusedAttnConfig &cfg) { - check_derived(cfg); - const auto cudnn_runtime_version = cudnnGetVersion(); - - b = static_cast(cfg.batch_size); - // keep original batch size because cu_seqlens are created with [b+1] shape - actual_b = b; - // The batch is bucketed unconditionally here, where the forward pass guards it: only the - // forward graph can be handed the user's cu_seqlens buffers directly, and it is their - // [actual_b+1] length that a quantized batch would overrun. The backward graph always reads - // converted seqlens out of our own workspace, so nothing here is sized by the true batch. - // make_cache_key() splits on the pass for this reason as well. - if ((cfg.is_ragged_q || cfg.is_ragged_kv) && cfg.uses_packed_ragged_graph) { - b = static_cast(cfg.bucketed_batch_size); - } - - // We choose between 32-bit and 64-bit offsets depending on need. - // This allows us to support older cuDNN runtimes gracefully. - ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; -} - -// The backward counterpart of build_sdpa_f16_fwd_graph; see there for why it constructs the graph -// and nothing else. -// -// Everything the graph's shape and topology depends on comes from `cfg` and `in`, so the build -// has one source of truth and cannot drift from the caller that will bind pointers to it. -static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig &cfg, - const F16BwdGraphInputs &in) { - const int64_t b = in.b; +// The backward counterpart of create_graph_f16_fwd; see there for why it constructs the graph and +// nothing else, and why the two direction-dependent dimensions are asked for rather than stored. +static SdpaF16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg) { + const GraphDims dims = graph_dims(cfg, Pass::Bwd); + const int64_t b = dims.batch_size; const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); const cudnn_frontend::DataType_t tensorType = @@ -710,7 +637,7 @@ static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig const auto cudnn_runtime_version = cudnnGetVersion(); const bool use_packed_ragged_graph = cfg.uses_packed_ragged_graph; const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = in.ragged_offset_type; + const DType ragged_offset_type = dims.ragged_offset_type; auto mha_graph = std::make_shared(); mha_graph->set_io_data_type(tensorType) @@ -940,12 +867,13 @@ static SdpaF16BwdGraphAndTensors build_sdpa_f16_bwd_graph(const FusedAttnConfig offset_s_tuple, dropout_tuple); } -// The backward counterpart of f16_fwd_cached_graph; see there. -static std::shared_ptr> f16_bwd_cached_graph( - const FusedAttnConfig &cfg, const F16BwdGraphInputs &in, cudnnHandle_t handle) { +// The backward counterpart of cache_graph_f16_fwd; see there. +static std::shared_ptr> cache_graph_f16_bwd( + const FusedAttnConfig &cfg, cudnnHandle_t handle) { static GraphCache cache; - return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::F16, Pass::Bwd, handle, - [&] { return build_sdpa_f16_bwd_graph(cfg, in); }); + check_derived(cfg); + return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Bwd), Backend::F16, Pass::Bwd, + handle, [&] { return create_graph_f16_bwd(cfg); }); } void fused_attn_arbitrary_seqlen_bwd_impl( @@ -958,12 +886,13 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cudnnHandle_t handle) { using namespace transformer_engine; - // Derived once and handed to the cache, which passes them to the graph build, so that the - // graph and the pointers bound to it below cannot be decided differently. - const F16BwdGraphInputs in(cfg); - const int64_t b = in.b; - const int64_t actual_b = in.actual_b; - const DType ragged_offset_type = in.ragged_offset_type; + // Asked with the same pass the graph was built with, so that the dimensions below and the ones + // the graph was built at cannot be decided differently. + const GraphDims dims = graph_dims(cfg, Pass::Bwd); + const int64_t b = dims.batch_size; + const DType ragged_offset_type = dims.ragged_offset_type; + // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by. + const int64_t actual_b = static_cast(cfg.batch_size); const bool use_ragged_stats = cfg.uses_ragged_stats; // Not const: bound into the variant pack by address as a pass-by-value graph input. @@ -976,13 +905,13 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const bool is_ragged_kv = cfg.is_ragged_kv; try { - auto cache_entry = f16_bwd_cached_graph(cfg, in, handle); + auto cache_entry = cache_graph_f16_bwd(cfg, handle); 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] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built(Backend::F16, Pass::Bwd, *cache_entry); + build_plans(Backend::F16, Pass::Bwd, *cache_entry); // 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. @@ -1322,48 +1251,24 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i // The question is answered by deriving the graph's inputs and building the graph, which is // where every rejection comes from -- there is no separate list of rules to keep in step with // the builder. The graph goes into the same cache the execution path reads, so the work is not -// thrown away and what was checked is what will run. It stops short of build_plans(), the +// thrown away and what was checked is what will run. It stops short of graph.build_plans(), the // expensive step, which the first execution of the graph does instead; see CachedGraph. // -// A refusal is cached too, so asking the same question twice costs one build rather than two; -// the second answer is the first one replayed. See GraphCache. +// A refusal, by contrast, is not cached: nothing is stored for a key cuDNN rejected, so asking the +// same question again pays for the build again. See lookup_or_cache_graph. // -// The copy below is made for the sake of one flag, which is not a redundant restatement of what -// the caller already asked for: make_cache_key() reads it to choose between the forward and the -// backward normalization, and one config can be probed in both directions -- the deprecated -// nvte_get_fused_attn_backend() leaves both check_for_*_support set, so both probes run off a -// single config. Each probe therefore states its own direction instead of inheriting it. +// The direction comes from which of these two functions was called, not from the config: a config +// arriving from a framework has both check_for_*_support set, so both probes run off a single +// config, and each has to name its own direction for the key and the graph to be the forward ones. std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - FusedAttnConfig graph_cfg = cfg; - graph_cfg.check_for_forward_support = true; - graph_cfg.check_for_backward_support = false; - - try { - const fused_attn::F16FwdGraphInputs in(graph_cfg); - fused_attn::f16_fwd_cached_graph(graph_cfg, in, handle); - return ""; - } catch (const std::exception &e) { - return fused_attn::refusal_reason(e, "is_supported_f16_fwd: rejected without a reason."); - } catch (...) { - return "is_supported_f16_fwd: unknown failure."; - } + return fused_attn::support_verdict("is_supported_f16_fwd", + [&] { fused_attn::cache_graph_f16_fwd(cfg, handle); }); } // The backward counterpart of is_supported_f16_fwd; see there. std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - FusedAttnConfig graph_cfg = cfg; - graph_cfg.check_for_forward_support = false; - graph_cfg.check_for_backward_support = true; - - try { - const fused_attn::F16BwdGraphInputs in(graph_cfg); - fused_attn::f16_bwd_cached_graph(graph_cfg, in, handle); - return ""; - } catch (const std::exception &e) { - return fused_attn::refusal_reason(e, "is_supported_f16_bwd: rejected without a reason."); - } catch (...) { - return "is_supported_f16_bwd: unknown failure."; - } + return fused_attn::support_verdict("is_supported_f16_bwd", + [&] { fused_attn::cache_graph_f16_bwd(cfg, 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 a4211275a3..8df1e3bfb2 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -21,9 +21,9 @@ using namespace transformer_engine; namespace fe = cudnn_frontend; // Every graph-cache event raised here names the build site it came from. This file is the fp8 -// backend throughout; only the pass differs between call sites. +// backend throughout; only the pass differs between call sites. Pass itself needs no +// using-declaration: it is fused_attn::Pass, since the config answers by direction too. using graph_cache_debug::Backend; -using graph_cache_debug::Pass; // fused attention FWD FP8 with FE 1.0+ using SdpaFp8FwdGraphAndTensors = @@ -49,60 +49,73 @@ using SdpaFp8FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The FP8 forward path's policy decisions: which quantization recipe the graph is built for, -// and whether cu_seqlens can be handed to cuDNN directly. Both decide which tensors the graph -// has, and so which pointers the variant pack has to bind -- the build and the execution cannot -// answer them differently, which is why they are derived once, here, for both. -struct Fp8FwdGraphInputs { - // Also where what FP8 cannot serve is rejected. Unlike the F16 path there is no bucketing to - // do, because FP8 has no ragged/THD support: the graph's shapes are exactly the config's. - explicit Fp8FwdGraphInputs(const FusedAttnConfig& cfg); - - bool is_delayed_scaling = false; - bool is_current_scaling = false; - bool is_mxfp8 = false; - bool use_cu_seqlens_directly = false; -}; - -Fp8FwdGraphInputs::Fp8FwdGraphInputs(const FusedAttnConfig& cfg) { +// The FP8 policy decisions a graph is built from, which the config cannot state on its own: which +// quantization recipe the graph quantizes for, whether cu_seqlens can be handed to cuDNN directly, +// and whether O arrives in F16. Each decides which tensors the graph has, and so which pointers +// the variant pack has to bind, so the build and the execution ask the same question here rather +// than each deciding for itself. Free functions for the reason check_derived() is one: they read +// public fields and compute, and belong to this backend rather than to the config. +// +// Unlike the F16 path there is no bucketing to do, because FP8 has no ragged/THD support: the +// graph's shapes are exactly the config's. + +// Which quantization recipe the graph is built for. One enum rather than the three mutually +// exclusive booleans it replaces, since exactly one recipe applies to a config and a triple leaves +// the other seven combinations expressible. +enum class Fp8Recipe { DelayedScaling, CurrentScaling, MxFp8 }; + +// The recipe `cfg` asks for, or a throw naming what FP8 cannot serve. The rejections are TE's own +// rather than cuDNN's -- bias, ALiBi and the recipe combinations -- and a support probe reports +// them the same way it reports a cuDNN refusal, as the reason the FP8 backend was not selected. +// +// The pass decides which tensor the recipe is read off: the forward graph writes O and the +// backward writes dQKV, and a run can quantize one without the other, so each pass reads the +// dtype of what it actually stores. +static Fp8Recipe fp8_recipe(const FusedAttnConfig& cfg, Pass pass) { check_derived(cfg); - const auto cudnn_runtime_version = cudnnGetVersion(); - const cudnn_frontend::DataType_t o_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + const cudnn_frontend::DataType_t out_type = + get_cudnn_fe_dtype(static_cast(pass == Pass::Fwd ? cfg.o_dtype : cfg.dqkv_dtype)); const NVTEScalingMode scaling_mode = cfg.scaling_mode; const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); const bool is_alibi = (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI); - const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); 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!"); - is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + const bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (out_type == cudnn_frontend::DataType_t::FP8_E4M3 || + out_type == cudnn_frontend::DataType_t::FP8_E5M2); + const bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (out_type == cudnn_frontend::DataType_t::HALF || + out_type == cudnn_frontend::DataType_t::BFLOAT16); + const bool is_mxfp8 = + (scaling_mode == NVTE_MXFP8_1D_SCALING) && (out_type == cudnn_frontend::DataType_t::HALF || + out_type == cudnn_frontend::DataType_t::BFLOAT16); NVTE_CHECK( is_delayed_scaling || is_current_scaling || is_mxfp8, "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); - NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, + NVTE_CHECK(!is_mxfp8 || cudnnGetVersion() >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike - // the F16 path, the FP8 path has no THD/ragged-offset support, so only the - // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the - // needed versions of cuDNN backend and frontend are higher than for F16.) - use_cu_seqlens_directly = + if (is_delayed_scaling) return Fp8Recipe::DelayedScaling; + if (is_current_scaling) return Fp8Recipe::CurrentScaling; + return Fp8Recipe::MxFp8; +} + +// Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative +// tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike +// the F16 path, the FP8 path has no THD/ragged-offset support, so only the +// cu_seqlens_to_actual_seqlens conversion applies here. Also note that the +// needed versions of cuDNN backend and frontend are higher than for F16, which is why this does +// not read cfg.uses_cu_seqlens_directly, the F16 path's answer to the same question.) +static bool fp8_uses_cu_seqlens_directly(const FusedAttnConfig& cfg) { + const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); + return // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). // Note: For the Python API, 1.27 is required. CUDNN_FRONTEND_VERSION >= 12600 && // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN // version, so we'll do the same. - (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && + (CUDNN_VERSION >= 92500 && cudnnGetVersion() >= 92500) && // 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 @@ -110,15 +123,25 @@ Fp8FwdGraphInputs::Fp8FwdGraphInputs(const FusedAttnConfig& cfg) { !is_dropout; } +// Whether O arrives in F16 rather than FP8, which is what decides if the backward graph has to +// descale it on the way in. Read off O for both passes, unlike the recipe. +static bool fp8_o_in_f16(const FusedAttnConfig& cfg) { + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + return o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16; +} + // Constructs the forward FP8 graph for one cache key, and only constructs it: whether cuDNN will -// run it is settled by the caller, in build_or_get_cached_graph(), which is also where the plan +// run it is settled by the caller, in lookup_or_cache_graph(), which is also where the plan // build eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every // call that does need one now sits on the other side of that boundary. // -// Everything the graph's shape and topology depends on comes from `cfg` and `in`, so the build -// has one source of truth and cannot drift from the caller that will bind pointers to it. -static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& cfg, - const Fp8FwdGraphInputs& in) { +// Everything the graph's shape and topology depends on comes from `cfg`, so the build has one +// source of truth and cannot drift from the caller that will bind pointers to it. The decisions +// the config cannot state itself are asked for with Pass::Fwd, the same way the code binding +// pointers to this graph asks. +static SdpaFp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t qkv_tensor_type = get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); @@ -149,10 +172,11 @@ static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& const bool is_padding = cfg.is_padding; const bool is_dropout = (is_training && dropout_probability != 0.0f); const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const bool is_delayed_scaling = in.is_delayed_scaling; - const bool is_current_scaling = in.is_current_scaling; - const bool is_mxfp8 = in.is_mxfp8; - const bool use_cu_seqlens_directly = in.use_cu_seqlens_directly; + const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Fwd); + const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; + const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; + const bool is_mxfp8 = recipe == Fp8Recipe::MxFp8; + const bool use_cu_seqlens_directly = fp8_uses_cu_seqlens_directly(cfg); auto mha_graph = std::make_shared(); mha_graph->set_io_data_type(qkv_tensor_type) @@ -398,11 +422,13 @@ static SdpaFp8FwdGraphAndTensors build_sdpa_fp8_fwd_graph(const FusedAttnConfig& // The FP8 forward graph cache and the only route to it. Both the execution path and the support // probe come through here, so a probe leaves behind exactly the entry a later execution finds. -static std::shared_ptr> fp8_fwd_cached_graph( - const FusedAttnConfig& cfg, const Fp8FwdGraphInputs& in, cudnnHandle_t handle) { +static std::shared_ptr> cache_graph_fp8_fwd( + const FusedAttnConfig& cfg, cudnnHandle_t handle) { static GraphCache cache; - return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::FP8, Pass::Fwd, handle, - [&] { return build_sdpa_fp8_fwd_graph(cfg, in); }); + // Asserted once here for both the key and the graph, which read the same derived fields. + check_derived(cfg); + return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Fwd), Backend::FP8, Pass::Fwd, + handle, [&] { return create_graph_fp8_fwd(cfg); }); } void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* devPtrK, @@ -415,13 +441,13 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - // Derived once and handed to the cache, which passes them to the graph build, so that the - // graph and the pointers bound to it below cannot be decided differently. Also where an - // unserviceable configuration is rejected. - const Fp8FwdGraphInputs in(cfg); - const bool is_delayed_scaling = in.is_delayed_scaling; - const bool is_current_scaling = in.is_current_scaling; - const bool use_cu_seqlens_directly = in.use_cu_seqlens_directly; + // Asked with the same pass the graph was built with, so that the tensors bound below and the + // ones the graph was built with cannot be decided differently. Also where an unserviceable + // configuration is rejected, ahead of the cache lookup. + const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Fwd); + const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; + const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; + const bool use_cu_seqlens_directly = fp8_uses_cu_seqlens_directly(cfg); const int64_t b = static_cast(cfg.batch_size); // Not const: bound into the variant pack by address as a pass-by-value graph input. @@ -432,13 +458,13 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); try { - auto cache_entry = fp8_fwd_cached_graph(cfg, in, handle); + auto cache_entry = cache_graph_fp8_fwd(cfg, handle); 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] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built(Backend::FP8, Pass::Fwd, *cache_entry); + build_plans(Backend::FP8, Pass::Fwd, *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -562,62 +588,15 @@ using SdpaFp8BwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// Builds the backward FP8 graph for one cache key, up to check_support() but not build_plans(); -// see CachedGraph for why the plan build is left to whoever executes the graph. +// Builds the backward FP8 graph for one cache key, up to check_support() but not +// graph.build_plans(); see CachedGraph for why the plan build is left to whoever runs the graph. // // Everything the graph's shape and topology depends on is re-derived from `cfg` here, so the // build has one source of truth for them. Unlike the F16 path, FP8 has no ragged/THD support, // so the shapes are exactly the config's and need no bucketing from the caller. -// The backward equivalent of Fp8FwdGraphInputs. The recipe is chosen from the dQKV dtype here -// rather than O's, since backward is what writes those. is_O_in_F16 additionally selects whether -// O has to be descaled on the way in. -struct Fp8BwdGraphInputs { - // The backward counterpart of Fp8FwdGraphInputs' constructor; see there for the rejections and - // for why the graph's shapes are simply the config's. - explicit Fp8BwdGraphInputs(const FusedAttnConfig& cfg); - - bool is_delayed_scaling = false; - bool is_current_scaling = false; - bool is_mxfp8 = false; - bool is_O_in_F16 = false; -}; - -Fp8BwdGraphInputs::Fp8BwdGraphInputs(const FusedAttnConfig& cfg) { - check_derived(cfg); - const auto cudnn_runtime_version = cudnnGetVersion(); - const cudnn_frontend::DataType_t o_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); - const cudnn_frontend::DataType_t dqkv_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.dqkv_dtype)); - const NVTEScalingMode scaling_mode = cfg.scaling_mode; - const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - const bool is_alibi = (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI); - - 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!"); - is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - NVTE_CHECK( - is_delayed_scaling || is_current_scaling || is_mxfp8, - "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); - NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, - "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - - is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); -} - -// The backward counterpart of build_sdpa_fp8_fwd_graph; see there for why it constructs the graph +// The backward counterpart of create_graph_fp8_fwd; see there for why it constructs the graph // and nothing else. -static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& cfg, - const Fp8BwdGraphInputs& in) { +static SdpaFp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t qkv_tensor_type = get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); @@ -655,10 +634,11 @@ static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& const bool is_padding = cfg.is_padding; const bool is_dropout = (dropout_probability != 0.0f); const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const bool is_delayed_scaling = in.is_delayed_scaling; - const bool is_current_scaling = in.is_current_scaling; - const bool is_mxfp8 = in.is_mxfp8; - const bool is_O_in_F16 = in.is_O_in_F16; + const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Bwd); + const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; + const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; + const bool is_mxfp8 = recipe == Fp8Recipe::MxFp8; + const bool is_O_in_F16 = fp8_o_in_f16(cfg); auto mha_graph = std::make_shared(); @@ -1033,12 +1013,13 @@ static SdpaFp8BwdGraphAndTensors build_sdpa_fp8_bwd_graph(const FusedAttnConfig& bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); } -// The backward counterpart of fp8_fwd_cached_graph; see there. -static std::shared_ptr> fp8_bwd_cached_graph( - const FusedAttnConfig& cfg, const Fp8BwdGraphInputs& in, cudnnHandle_t handle) { +// The backward counterpart of cache_graph_fp8_fwd; see there. +static std::shared_ptr> cache_graph_fp8_bwd( + const FusedAttnConfig& cfg, cudnnHandle_t handle) { static GraphCache cache; - return build_or_get_cached_graph(cache, cfg.make_cache_key(), Backend::FP8, Pass::Bwd, handle, - [&] { return build_sdpa_fp8_bwd_graph(cfg, in); }); + check_derived(cfg); + return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Bwd), Backend::FP8, Pass::Bwd, + handle, [&] { return create_graph_fp8_bwd(cfg); }); } void fused_attn_fp8_bwd_impl( @@ -1055,14 +1036,14 @@ void fused_attn_fp8_bwd_impl( cudnnHandle_t handle) { using namespace transformer_engine; - // Derived once and handed to the cache, which passes them to the graph build, so that the - // graph and the pointers bound to it below cannot be decided differently. Also where an - // unserviceable configuration is rejected. - const Fp8BwdGraphInputs in(cfg); - const bool is_delayed_scaling = in.is_delayed_scaling; - const bool is_current_scaling = in.is_current_scaling; - const bool is_mxfp8 = in.is_mxfp8; - const bool is_O_in_F16 = in.is_O_in_F16; + // Asked with the same pass the graph was built with, so that the tensors bound below and the + // ones the graph was built with cannot be decided differently. Also where an unserviceable + // configuration is rejected, ahead of the cache lookup. + const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Bwd); + const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; + const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; + const bool is_mxfp8 = recipe == Fp8Recipe::MxFp8; + const bool is_O_in_F16 = fp8_o_in_f16(cfg); const int64_t b = static_cast(cfg.batch_size); const int64_t h = static_cast(cfg.num_attn_heads); @@ -1074,7 +1055,7 @@ void fused_attn_fp8_bwd_impl( const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); try { - auto cache_entry = fp8_bwd_cached_graph(cfg, in, handle); + auto cache_entry = cache_graph_fp8_bwd(cfg, handle); 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, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, @@ -1082,7 +1063,7 @@ void fused_attn_fp8_bwd_impl( dropout_seed, dropout_offset] = cache_entry->tensors; // This graph is going to be used, so finish the build the cache deferred. - ensure_plans_built(Backend::FP8, Pass::Bwd, *cache_entry); + build_plans(Backend::FP8, Pass::Bwd, *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1411,46 +1392,29 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const } } -// Whether cuDNN can run the FP8 forward graph this config asks for: the empty string if it can, -// otherwise cuDNN's own account of why not, which the backend selector reports to the caller. +// Whether the FP8 forward graph this config asks for can run: the empty string if it can, +// otherwise the account of why not, which the backend selector reports to the caller. // -// The question is answered by deriving the graph's inputs and building the graph, which is -// where every rejection comes from -- there is no separate list of rules to keep in step with -// the builder. The graph goes into the same cache the execution path reads, so the work is not -// thrown away and what was checked is what will run. It stops short of build_plans(), the -// expensive step, which the first execution of the graph does instead; see CachedGraph. +// The question is answered by building the graph, which is where every rejection comes from -- +// there is no separate list of rules to keep in step with the builder. The graph goes into the same +// cache the execution path reads, so the work is not thrown away and what was checked is what will +// run. It stops short of graph.build_plans(), the expensive step, which the first execution of the +// graph does instead; see CachedGraph. // -// The copy below is made for the sake of one flag; see is_supported_f16_fwd for why that -// assignment is direction selection rather than a restatement of the caller's request. +// Unlike the F16 path, some of the rejections here are TE's own rather than cuDNN's: fp8_recipe() +// throws for bias, ALiBi and the recipe combinations FP8 does not serve, from inside the build. +// They read the same to the selector, which wants a reason and does not care whose rule it was. +// +// The direction comes from which of these two functions was called; see is_supported_f16_fwd. std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { - FusedAttnConfig graph_cfg = cfg; - graph_cfg.check_for_forward_support = true; - - try { - const fused_attn::Fp8FwdGraphInputs in(graph_cfg); - fused_attn::fp8_fwd_cached_graph(graph_cfg, in, handle); - return ""; - } catch (const std::exception& e) { - return fused_attn::refusal_reason(e, "is_supported_fp8_fwd: rejected without a reason."); - } catch (...) { - return "is_supported_fp8_fwd: unknown failure."; - } + return fused_attn::support_verdict("is_supported_fp8_fwd", + [&] { fused_attn::cache_graph_fp8_fwd(cfg, handle); }); } // The backward counterpart of is_supported_fp8_fwd; see there. std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { - FusedAttnConfig graph_cfg = cfg; - graph_cfg.check_for_forward_support = false; - - try { - const fused_attn::Fp8BwdGraphInputs in(graph_cfg); - fused_attn::fp8_bwd_cached_graph(graph_cfg, in, handle); - return ""; - } catch (const std::exception& e) { - return fused_attn::refusal_reason(e, "is_supported_fp8_bwd: rejected without a reason."); - } catch (...) { - return "is_supported_fp8_bwd: unknown failure."; - } + return fused_attn::support_verdict("is_supported_fp8_bwd", + [&] { fused_attn::cache_graph_fp8_bwd(cfg, handle); }); } } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index 7a226c71fb..d96d7166ca 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -8,31 +8,30 @@ // The fused-attention graph cache: what a cache entry is, how one is looked up // or built, and the frontend calls that make a constructed graph usable. // -// Each of the four build sites (f16 and fp8, forward and backward) differs only -// in how it constructs its graph and which tensors it hands back. Everything -// after that -- the lookup, the locking, the once-per-entry plan build, the -// support check, and the remembering of what cuDNN refused -- is the same at all -// four, and lives here so it has one definition rather than four copies to keep -// in step. +// The four build sites (f16 and fp8, forward and backward) differ only in how +// they construct their graph and which tensors they hand back. Everything after +// that -- the lookup, the locking, the support check, the once-per-entry plan +// build -- is shared, and lives here rather than in four copies. // -// The five frontend calls a graph goes through, and which caller pays for each: +// The five frontend calls a graph goes through, and which of our functions pays for +// each. The frontend's are written graph.*, since that is how they are invoked and +// since two of them share a name with ours: // // on a miss, either caller: -// validate() -> build_operation_graph() -> create_execution_plans(HeurMode_t::A) -// -> check_support() validate_and_check_support() +// graph.validate() -> graph.build_operation_graph() +// -> graph.create_execution_plans(HeurMode_t::A) -> graph.check_support() +// all four via query_support() // the execution path only: -// build_plans() ensure_plans_built(), once per entry, the kernel compilation -// execute() every call, with its variant pack built in a local +// graph.build_plans() build_plans(), once per entry, the kernel compilation +// graph.execute() every call, with its variant pack built in a local // -// This header is deliberately not part of utils.h: it needs the cuDNN frontend, -// and utils.h is included by translation units (utils.cu) that otherwise do not. +// Not part of utils.h: this needs the cuDNN frontend, and utils.h is included by +// translation units (utils.cu) that otherwise do not. // ============================================================================ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ -#include -#include #include #include #include @@ -50,76 +49,44 @@ namespace transformer_engine { namespace fused_attn { -// cuDNN's refusal to run a graph, as opposed to a failure to try. The distinction is what makes -// the negative cache in build_or_get_cached_graph() safe: a refusal is a verdict on the -// configuration and reproducible for a given key, so it can be remembered and replayed, whereas -// a failure that came from the machine's state at that moment (an allocation that did not fit, a -// CUDA error left behind by unrelated work) could well succeed on the next attempt and must not -// be turned into a permanent answer. Only a frontend call that returned one of the codes -// is_unsupported_verdict() names raises this; every other failure keeps its ordinary type and is -// re-attempted the next time the key comes around. -struct UnsupportedGraph : public std::runtime_error { - explicit UnsupportedGraph(const std::string &reason) : std::runtime_error(reason) {} -}; - -// Whether a frontend error code is a verdict on the graph rather than a report of something -// that went wrong on the way to reaching one. +// The verdict an is_supported_* helper reports for `probe`: the empty string if it completes, +// otherwise cuDNN's own account of the refusal. Support is discovered by building the graph, so a +// probe is one call and everything else is what to do with a failure; this is that, once, for all +// four helpers. // -// cudnn-frontend distinguishes the two, and the negative cache is only sound for the first. -// Three codes are verdicts: +// Refusals and failures on the way to a verdict read alike, because CUDNN_BACKEND_API_FAILED -- +// raised for any non-success cudnnStatus_t -- cannot separate CUDNN_STATUS_NOT_SUPPORTED from +// CUDNN_STATUS_ALLOC_FAILED. Either way this backend cannot serve this call, and either way what +// the caller wants is the message. // -// GRAPH_NOT_SUPPORTED is what the frontend's own support surface returns, from validate(). -// Nearly every rule it checks by hand -- the architecture gates, the head-dim limits, the -// version-specific workarounds -- reports itself this way. -// GRAPH_EXECUTION_PLAN_CREATION_FAILED is what check_support() returns when no engine config -// cuDNN's heuristics offered can run the graph. This is the verdict for everything the -// frontend does not rule on itself and defers to the backend, so leaving it out would -// exclude most of what a support probe actually discovers. -// UNSUPPORTED_GRAPH_FORMAT is a verdict by name and costs nothing to accept, though no -// frontend release we build against returns it. -// -// All three are properties of the key and will be just as true the next time it is asked. Every -// other code -- CUDNN_BACKEND_API_FAILED, CUDA_API_FAILED, HEURISTIC_QUERY_FAILED, HANDLE_ERROR, -// INVALID_CUDA_DEVICE and the rest -- describes the process at that moment: an OOM under memory -// pressure, a sticky CUDA error left by unrelated work, a handle on the wrong device. -// CUDNN_BACKEND_API_FAILED is the one to be careful about, since the frontend raises it for any -// non-success cudnnStatus_t and so cannot tell CUDNN_STATUS_ALLOC_FAILED from -// CUDNN_STATUS_NOT_SUPPORTED; caching it would let a moment of memory pressure blacklist a -// configuration that is genuinely supported, for the life of the process, and the wider the -// cache's reach the worse that gets. So those are raised as ordinary errors, which leave nothing -// behind and are retried when the key next comes around. Either way cuDNN's own message reaches -// the caller; only whether it is remembered differs. -inline bool is_unsupported_verdict(cudnn_frontend::error_code_t code) { - return code == cudnn_frontend::error_code_t::GRAPH_NOT_SUPPORTED || - code == cudnn_frontend::error_code_t::GRAPH_EXECUTION_PLAN_CREATION_FAILED || - code == cudnn_frontend::error_code_t::UNSUPPORTED_GRAPH_FORMAT; -} - -// The reason string an is_supported_* helper reports for `e`: its message, or `fallback` if it -// has none. Those helpers signal support by returning the empty string, so a refusal that -// arrives without an explanation would be read as an endorsement and the caller would go on to -// run a graph cuDNN has just declined. Nothing raised through NVTE_ERROR can be empty, since it -// prefixes file and line, but that is a property of our macros rather than of every exception -// that can reach a catch clause, and it is not what the contract should rest on. -inline std::string refusal_reason(const std::exception &e, const char *fallback) { - const char *what = e.what(); - return (what != nullptr && what[0] != '\0') ? std::string(what) : std::string(fallback); +// `what` names the probe and is used only when a failure carried no message of its own: support is +// signalled by returning the empty string, so an empty refusal would read as an endorsement. +template +std::string support_verdict(const char *what, ProbeFn &&probe) { + try { + probe(); + return ""; + } catch (const std::exception &e) { + const char *reason = e.what(); + if (reason != nullptr && reason[0] != '\0') return reason; + return std::string(what) + ": rejected without a reason."; + } catch (...) { + return std::string(what) + ": unknown failure."; + } } // A graph in the cache, plus the tensor attributes needed to bind runtime pointers to it. // -// Entries are built only as far as check_support(), which is all it takes to decide whether -// a configuration is supported. build_plans() is the kernel-compilation step and the most -// expensive of the five frontend calls, so a support query stops short of it: the query never -// executes the graph, and many of the keys it builds are never executed by anything. The -// execution path finishes the build instead, the first time the graph is needed to run. +// Entries are built only as far as check_support(), which is all it takes to decide whether a +// configuration is supported. graph.build_plans() -- the kernel compilation, and the most expensive +// of the five frontend calls -- is left to the execution path, since a support query never runs the +// graph and many of the keys it builds are never run by anything. // -// plans_built guards that completion. It has to happen exactly once per entry, because the -// cached graph is shared across threads and build_plans() mutates it in place -- two threads -// reaching the same unfinished entry must not both build it. Keeping the flag inside the entry -// keeps it from drifting away from the graph it describes, and leaves unrelated keys free to -// build concurrently. A build that throws leaves the flag unset, so a later call retries -// rather than executing a graph with no plans. +// plans_built guards that completion, which has to happen exactly once per entry: the entry is +// shared across threads and graph.build_plans() mutates it in place. Keeping the flag in the entry +// keeps it with the graph it describes and leaves unrelated keys free to build concurrently. A +// build that throws leaves it unset, so a later call retries rather than executing a graph with no +// plans. template struct CachedGraph { explicit CachedGraph(GraphAndTensors tensors) : tensors(std::move(tensors)) {} @@ -128,130 +95,57 @@ struct CachedGraph { std::once_flag plans_built; }; -// One build site's cache. Process-wide rather than per-thread so that a graph is reused -// across threads instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a -// shared plan, and the frontend's execute() builds its variant pack in a local rather than in -// the graph, so it does not write to the shared object. No particular frontend version is -// relied on for that -- it has held for far longer than the >= 1.25.0 the build requirements -// ask for, which is there for unrelated features. +// One build site's cache. Process-wide rather than per-thread so a graph is reused across threads +// instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a shared plan, and the +// frontend's execute() builds its variant pack in a local rather than in the graph, so it does not +// write to the shared object. // // What lets one cache serve every thread is an asymmetry between the two objects a call needs. A -// cuDNN handle is per-thread mutable session state: it carries the stream that execute() launches -// on, so each thread holds its own rather than racing to set that on a shared one. A graph and -// its plans are the opposite -- compiled artifacts, built for the properties of a device and -// bound to the device they were finalized against, with nothing in them belonging to the thread -// that did the building. So the cache can be keyed by device and shared by all threads, which is -// why make_cache_key() stamps device_id and nothing thread-shaped. ensure_plans_built() covers -// what that costs at the seam, where the thread that finishes a build is often not the thread -// that started it. +// cuDNN handle is per-thread mutable session state: it carries the stream execute() launches on, so +// each thread holds its own. A graph and its plans are the opposite -- compiled artifacts, bound to +// the device they were finalized against, with nothing in them belonging to the building thread. So +// the key stamps device_id and nothing thread-shaped (see make_cache_key), and build_plans() below +// covers the seam where the thread that finishes a build is not the one that started it. // -// Refusals are cached alongside the graphs, under the same keys and the same lock. A support -// query for an unsupported configuration is otherwise the most expensive thing this cache sees: -// it builds the whole graph, spends the four frontend calls, and throws the result away, and it -// does so again on every query, because a rejection left nothing behind to find. `unsupported` -// is what it leaves behind -- cuDNN's own account of the refusal, which is the entire useful -// output of a failed query, so nothing is lost by answering from it. Reasons are short strings -// and there is one per refused key, so this grows far slower than the graphs beside it. -// Holding the lock and the maps together is also what fixes their relative lifetimes. Members -// are destroyed in reverse declaration order, so the mutex is declared first to be destroyed -// last: the maps go while their guard is still valid, rather than the other way round. Declaring -// a cache and its lock as two separate objects leaves that ordering to whoever writes the next -// one; declaring them here settles it once. +// The mutex is declared first so that it is destroyed last -- members go in reverse declaration +// order, so the map is destroyed while its guard is still valid. Declaring the two together settles +// that rather than leaving it to whoever writes the next cache. // -// Both maps are bounded; see kCacheCapacity. `last_used` is what makes the bound an LRU rather -// than an arbitrary cull: it is stamped from `clock` on every insertion and every hit, so the -// entry with the smallest value is the one that has gone longest without being asked for. The -// clock is an ordinary member rather than an atomic because it is only ever touched under -// `mutex`, alongside the maps it orders. +// The map is unbounded. Only executed graphs hold anything substantial -- an entry that stopped at +// check_support() has no compiled kernels behind it, and none hold a workspace, which the caller +// allocates per call -- and a model reuses a handful of configurations, so any bound worth setting +// would sit far above what real work reaches. A workload that does sweep shapes, such as a suite +// enumerating them, holds every graph for the life of the process; `miss` climbing without settling +// is what that looks like, and is the case for bringing a bound back. template struct GraphCache { - struct Slot { - std::shared_ptr> entry; - uint64_t last_used; - }; - struct Refusal { - std::string reason; - uint64_t last_used; - }; - std::mutex mutex; // guards everything below - uint64_t clock = 0; - std::map supported; - std::map unsupported; + std::map>> entries; }; -// The ceiling on entries in one of the maps of one build site's cache. -// -// Sized to be out of the way of real work rather than to be tight. A training step reuses a -// handful of configurations and an inference server with bucketed sequence lengths tens of them, -// so a hundred is already more shape diversity than a model exhibits. What the ceiling is for is -// the case where the key space is effectively unbounded -- a test suite sweeping shapes, or a -// serving workload that keys on something that never repeats -- where an unbounded cache is a -// slow leak of cuDNN graphs and their execution plans for the life of the process. -// -// Hard-coded rather than configurable, because nothing has yet needed a different number: the -// workloads that fit under it never notice the ceiling, and the ones that do not are better -// served by rebuilding a graph than by holding thousands. An environment variable can come back -// if a workload turns up that wants to trade the memory for the rebuilds. -constexpr size_t kCacheCapacity = 100; - -// Make room in `entries` for one more, by dropping the least recently used until there is. -// Call under the cache's lock. -// -// Evicting a graph does not invalidate one that is in use. build_or_get_cached_graph() hands -// back a shared_ptr, so a thread that is executing an entry holds it alive regardless of what -// the map does; erasing here drops the cache's reference and nothing else. The scan is linear, -// but it runs only when the cache is full, and comparing a hundred integers is nothing beside -// the graph build it is making room for. -template -void evict_to_fit(Map &entries) { - while (entries.size() >= kCacheCapacity) { - auto oldest = entries.begin(); - for (auto it = entries.begin(); it != entries.end(); ++it) { - if (it->second.last_used < oldest->second.last_used) oldest = it; - } - entries.erase(oldest); - } -} - // Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: -// validate, build_operation_graph, create_execution_plans, check_support. The sequence is -// identical for both passes and both backends, so it is defined once here; `backend` and `pass` -// only name the build site whose stage timers the calls are attributed to. -// -// Support is reported by throwing rather than by a return value. NVTE_CHECK_CUDNN_FE raises -// an exception carrying cuDNN's own explanation of the rejection, and that text is what the -// is_supported_* helpers return as the reason a backend was refused -- so a bool here would -// discard the one thing a support probe exists to produce. Callers that are about to execute -// the graph want the throw as well, since there is nothing useful to do with an unsupported -// graph but fail. +// validate, build_operation_graph, create_execution_plans, check_support. Identical for both passes +// and both backends, so it is defined once here; `backend` and `pass` only name the build site the +// stage timers attribute the calls to. // -// A failure is raised as UnsupportedGraph only when the frontend's own error code says the graph -// was adjudicated and refused; see is_unsupported_verdict(). Anything else these calls can report -// is a failure to reach a verdict and is raised through NVTE_ERROR, so it is not remembered. -// Classifying on the code rather than on which call failed is what keeps that honest: all four of -// these calls can fail for environmental reasons too -- build_operation_graph() and -// create_execution_plans() both talk to the cuDNN backend -- so their position in the sequence -// says nothing about whether the failure was about the configuration. +// Reports by throwing, and the throw carries cuDNN's message alone. That message is what the +// is_supported_* helpers return as the reason a backend was refused, so a bool would discard the +// one thing a support probe exists to produce -- and NVTE_ERROR would wrap it in the file, line +// and advice of an internal failure, which a backend refused for a plain reason is not. One kind +// of throw for every failure; see support_verdict() for why that distinction is not drawn. // -// build_plans() and execute() sit outside this function entirely: they commit real resources, and -// build_plans() belongs to whoever executes the graph, once, the first time it is needed. See -// CachedGraph. -inline void validate_and_check_support(graph_cache_debug::Backend backend, - graph_cache_debug::Pass pass, - cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { +// graph.build_plans() and graph.execute() sit outside this function: they commit real resources, +// and the plan build belongs to whoever executes the graph, once. See CachedGraph. +inline void query_support(graph_cache_debug::Backend backend, Pass pass, + cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { cudnn_frontend::error_t error; - graph_cache_debug::timer(backend, pass, stage, [&] { error = call(); }); + graph_cache_debug::record_time(backend, pass, stage, [&] { error = call(); }); if (error.is_good()) return; // cuDNN normally explains itself; fall back to the call's name so that a refusal can never // arrive as an empty string, which the is_supported_* helpers would read as an endorsement. - const std::string reason = - error.err_msg.empty() ? std::string(call_name) + " failed." : error.err_msg; - if (is_unsupported_verdict(error.code)) throw UnsupportedGraph(reason); - NVTE_ERROR("cuDNN Error in ", call_name, ": ", reason, - " For more information, enable cuDNN error logging by setting CUDNN_LOGERR_DBG=1 " - "and CUDNN_LOGDEST_DBG=stderr in the environment."); + throw std::runtime_error(error.err_msg.empty() ? std::string(call_name) + " failed." + : error.err_msg); }; run(graph_cache_debug::BuildStage::Validate, "validate", [&] { return graph.validate(); }); @@ -263,157 +157,107 @@ inline void validate_and_check_support(graph_cache_debug::Backend backend, [&] { return graph.check_support(); }); } -// The cached entry for `key`, building and inserting it via `build` if absent. Throws -// UnsupportedGraph if cuDNN refuses the graph -- this time or on an earlier call, the two being -// indistinguishable to the caller by design. +// The cached entry for `key`, building and inserting it via `build` if absent. Throws if cuDNN +// refuses the graph, and remembers nothing when it does, so the next query for a refused key builds +// it again and is refused again. The frameworks only re-enter the selector when the attention +// configuration changes (in PyTorch, _attention_backends caches the choice), so a settled run pays +// for a refusal once; a suite that enumerates configurations pays each time it comes back around, +// which is the case a map of remembered refusals would serve. // -// `build` only constructs a graph; this is what puts it through validate_and_check_support(), so -// the entries in the cache are exactly the graphs cuDNN has agreed to run. Those four calls sit -// on the miss path because they are part of building an entry rather than reading one: repeating -// them on a hit would redo the operation graph and the plan search for a graph that has already -// been through both. +// `build` only constructs a graph; this is what puts it through query_support(), so the entries in +// the cache are exactly the graphs cuDNN has agreed to run. Those calls belong to building an entry +// rather than reading one, which is why a hit skips them. // -// `key` must be a normalized key -- FusedAttnConfig::make_cache_key()'s output -- and not a -// raw execution config. Two configs that differ only in a field no graph reads (attn_scale, -// say) have to reach the same entry, which is what normalization is for; passing the raw -// config instead silently multiplies the cache by fields the graph never consumes. +// `key` must be make_cache_key(pass)'s output, for the same `pass` given here, and not a raw +// execution config: two configs differing only in a field no graph reads (attn_scale, say) have to +// reach the same entry, and passing the raw config silently multiplies the cache by fields the +// graph never consumes. // -// Only the map operations are locked, not `build`. A graph build is the expensive part and -// holding the lock across it would serialize builds of unrelated keys, so two threads racing -// on the same key may both build. That is a wasted build, not a correctness problem: the -// loser drops its own graph and takes the winner's, so every caller of a given key gets one -// shared entry and the once-flag inside it still governs the plan build. Both threads record -// their own lookup, so the wasted build shows up in diagnostics as two MISS lines carrying the -// same key and a build_graph count above the number of distinct keys, rather than as anything -// missing. The same race on a refused key is equally harmless, both threads storing the same -// reason. +// Only the map operations are locked, not `build`, so builds of unrelated keys proceed concurrently +// and two threads racing on one key may both build. That is wasted work rather than a correctness +// problem -- the loser drops its graph and takes the winner's entry, whose once_flag still governs +// the plan build -- and it reads in diagnostics as two MISS lines with the same key. // // lock cache.mutex -// supported[key]? found -> last_used = ++clock, copy the shared_ptr -// unsupported[key]? found -> last_used = ++clock, copy the reason +// entries[key]? found -> copy the shared_ptr // unlock -// record_cache_lookup(HIT | UNSUPPORTED | MISS) +// record_cache_lookup(HIT | MISS) // -// HIT -> return the entry -// UNSUPPORTED -> throw UnsupportedGraph(the remembered reason) -// MISS -> build() outside the lock, so builds of unrelated -// validate_and_check_support() keys proceed concurrently -// ok -> lock, evict_to_fit(supported), insert stamped ++clock, unlock, -// return the inserted entry, which on a lost race is the winner's -// verdict -> lock, evict_to_fit(unsupported), insert the reason, unlock, -// rethrow -// other -> NVTE_ERROR: nothing remembered, retried when the key returns +// HIT -> return the entry +// MISS -> build() outside the lock, so builds of unrelated +// query_support() keys proceed concurrently +// ok -> lock, insert, unlock, return the inserted entry, which on a lost race +// is the winner's +// throw -> propagates; nothing is stored, so the key is built again if it comes back template -std::shared_ptr> build_or_get_cached_graph( +std::shared_ptr> lookup_or_cache_graph( GraphCache &cache, const FusedAttnConfig &key, - graph_cache_debug::Backend backend, graph_cache_debug::Pass pass, cudnnHandle_t handle, - BuildFn &&build) { + graph_cache_debug::Backend backend, Pass pass, cudnnHandle_t handle, BuildFn &&build) { using Entry = CachedGraph; - using Slot = typename GraphCache::Slot; - using Refusal = typename GraphCache::Refusal; + using graph_cache_debug::LookupResult; std::shared_ptr cached; - bool refused = false; - std::string reason; { std::lock_guard lock(cache.mutex); - auto it = cache.supported.find(key); - if (it != cache.supported.end()) { - it->second.last_used = ++cache.clock; - cached = it->second.entry; - } else { - auto refusal = cache.unsupported.find(key); - refused = (refusal != cache.unsupported.end()); - if (refused) { - refusal->second.last_used = ++cache.clock; - reason = refusal->second.reason; - } - } - } - using graph_cache_debug::LookupResult; - LookupResult outcome = LookupResult::Miss; - if (cached != nullptr) { - outcome = LookupResult::Hit; - } else if (refused) { - outcome = LookupResult::Unsupported; + auto it = cache.entries.find(key); + if (it != cache.entries.end()) cached = it->second; } - // Recorded after the lock is dropped, so that writing a trace line cannot hold up threads - // querying other keys. The counters are exact, but two lookups that raced on the lock can be - // recorded in the opposite order, so read a level-2 trace as the set of lookups that happened - // rather than as the sequence they happened in. - graph_cache_debug::record_cache_lookup(backend, pass, outcome, key); - + // Recorded after the lock is dropped, so writing a trace line cannot hold up threads querying + // other keys. The counters stay exact, but two lookups that raced can be recorded in the opposite + // order, so a level-2 trace is the set of lookups that happened, not their sequence. + graph_cache_debug::record_cache_lookup( + backend, pass, cached != nullptr ? LookupResult::Hit : LookupResult::Miss, key); if (cached != nullptr) return cached; - // Raised rather than returned so that a replayed refusal is the same event as a fresh one: - // every caller already has to handle the build refusing, and none of them would have anything - // else to do with a second, quieter way of saying so. - if (refused) throw UnsupportedGraph(reason); - std::shared_ptr entry; - try { - entry = std::make_shared(build()); - // Every site's tensor tuple leads with its graph, which is the one thing all four have in - // common and the only element this needs. A tuple that stopped leading with it would fail to - // compile here rather than quietly validate the wrong object. - validate_and_check_support(backend, pass, *std::get<0>(entry->tensors), handle); - } catch (const UnsupportedGraph &e) { - { - std::lock_guard lock(cache.mutex); - evict_to_fit(cache.unsupported); - cache.unsupported.insert({key, Refusal{e.what(), ++cache.clock}}); - } - graph_cache_debug::record_unsupported(backend, pass); - throw; - } - graph_cache_debug::record_graph_built(backend, pass); + // A failure propagates with cuDNN's message and leaves nothing behind. It raised a MISS and no + // CREATE_GRAPH, which is what makes miss - create_graph the count of builds that ended this way. + auto entry = std::make_shared(build()); + // Every site's tensor tuple leads with its graph, the one element this needs. A tuple ordered + // otherwise would fail to compile rather than quietly validate the wrong object. + query_support(backend, pass, *std::get<0>(entry->tensors), handle); + graph_cache_debug::record_graph_created(backend, pass); { std::lock_guard lock(cache.mutex); - evict_to_fit(cache.supported); - // On a losing race the insert does nothing: the temporary Slot is destroyed with the graph - // this thread built, and what comes back is the winner's entry. - auto inserted = cache.supported.insert({key, Slot{std::move(entry), ++cache.clock}}); - return inserted.first->second.entry; + // On a losing race the insert does nothing: the shared_ptr this thread built is dropped with + // its graph, and what comes back is the winner's entry. + auto inserted = cache.entries.insert({key, std::move(entry)}); + return inserted.first->second; } } -// Runs the plan build that build_or_get_cached_graph() left undone, once per entry. -// -// Call this only when the graph is about to be executed, which is why it is a separate step -// rather than the tail of the lookup: a support query builds entries that nothing ever runs, and -// kernel compilation is the most expensive of the five frontend calls, so a query that paid for -// it would be paying for nothing. See CachedGraph for why the flag lives inside the entry and -// what a throw here leaves behind. -// -// Splitting the build in two means the thread that finishes it is often not the thread that -// started it -- a sizing call on one thread caches the graph, and an autograd thread is the first -// to need it to run. Four facts make that safe, and only the first is visible here. -// -// build_plans() takes no handle. The overload that accepts one ignores it -- its body is -// `(void)handle;` -- and the build works from the operation graph descriptor and the device -// properties instead, which is how deviceless ahead-of-time compilation builds plans with no -// handle at all. Unlike the plan sharing described on GraphCache, this does lean on the >= 1.25.0 -// frontend the build requires: it is where the handle-free overload arrived. Calling it means a -// plan build cannot reach for the handle of a thread that has since exited. -// -// The handle from the build does outlive the build, held by the operation graph descriptor that -// build_operation_graph(handle) finalized against it. It stays a valid object only because TE -// never destroys cuDNN handles: cudnnExecutionPlanManager leaves HandleManager's Destroy -// parameter at its nullptr default, so handles leak by design, one per thread per device. +// Runs graph.build_plans(), the plan build that lookup_or_cache_graph() left undone, once per +// entry. Named for the frontend call it wraps; the once-per-entry part is the whole reason it is a +// function rather than that call. // -// That descriptor was finalized for the device of the handle that built it, which is why the -// cache key carries device_id (see FusedAttnConfig::make_cache_key). Without it a thread could -// build plans, and compile kernels, from a descriptor belonging to another device. +// Call only when the graph is about to be executed, which is why this is a separate step rather +// than the tail of the lookup: a support query builds entries nothing ever runs, and kernel +// compilation is the most expensive of the five frontend calls. See CachedGraph for why the flag +// lives inside the entry and what a throw here leaves behind. // -// Execution stays clear of all of it: execute() is called with the running thread's own handle, -// so a handle is never used by two threads at once, which is what cuDNN asks in return for -// letting them share the plan. +// Splitting the build in two means the thread that finishes it is often not the thread that started +// it -- a sizing call caches the graph, and an autograd thread is the first to need it to run. Four +// facts make that safe, and only the first is visible here: +// - graph.build_plans() takes no handle. The overload that accepts one ignores it (its body is +// `(void)handle;`), working from the operation graph descriptor and the device properties +// instead, which is how deviceless ahead-of-time compilation builds plans with no handle at +// all. Unlike the plan sharing on GraphCache, this does lean on the >= 1.25.0 frontend the +// build requires: it is where the handle-free overload arrived. +// - The handle that built the operation graph outlives the build, held by the descriptor +// graph.build_operation_graph(handle) finalized against it, and stays valid only because TE +// never destroys cuDNN handles: cudnnExecutionPlanManager leaves HandleManager's Destroy +// parameter at its nullptr default, so handles leak by design, one per thread per device. +// - That descriptor was finalized for one device, which is why the cache key carries device_id +// (see make_cache_key). Without it a thread could compile kernels from another device's +// descriptor. +// - graph.execute() is called with the running thread's own handle, so a handle is never used by +// two threads at once, which is what cuDNN asks in return for letting them share a plan. template -void ensure_plans_built(graph_cache_debug::Backend backend, graph_cache_debug::Pass pass, - CachedGraph &entry) { +void build_plans(graph_cache_debug::Backend backend, Pass pass, + CachedGraph &entry) { std::call_once(entry.plans_built, [&] { cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.tensors); - graph_cache_debug::timer(backend, pass, graph_cache_debug::BuildStage::BuildPlans, - [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); + graph_cache_debug::record_time(backend, pass, graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); graph_cache_debug::record_plans_built(backend, pass); }); } diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 4222843ce9..c19a2902b0 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -7,75 +7,39 @@ // ============================================================================ // Fused-attention graph cache diagnostics. // -// Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG. Two verbosity levels: -// =1 (events) : low volume. Cache event counters, a BUILD_GRAPH and a BUILD_PLANS line -// per build, an UNSUPPORTED line per configuration cuDNN refuses, and the -// end-of-run SUMMARY (per backend and per thread, plus a row across the -// backends when a run used more than one) with stage timings. Each of these -// fires once per distinct cache key, which is what keeps the volume low, and -// is enough to diagnose redundant rebuilds and profile build cost. -// =2 (trace) : high volume. Additionally emits a per-lookup HIT/MISS/UNSUPPORTED line -// with the full shorthand cache key and a per-execution EXEC line. Use only -// when you need to see *which* shapes are hitting/missing -- these fire on -// every cache lookup and execution, so at suite scale they add I/O and -// serialize threads on the stderr lock. No timed region writes to stderr, so -// the stage timings stay sound, but they are measured under more contention -// than at level 1 and read a little high. +// Enable with NVTE_FUSED_ATTN_CACHE_DEBUG=[:]. The output format, how +// to read it and the rank suffix are documented for users in docs/envvars.rst; what +// follows is what maintaining this file needs. // -// Every line names the build site behind it, "f16" or "fp8" followed by the pass, and the -// counters it carries belong to that backend alone -- the two keep separate columns, so a -// process that drives both can still say which of them built what. Every event name is also -// the counter column it increments, so a line and the totals beside it read with one -// vocabulary. UNSUPPORTED names both a level-1 event and a level-2 lookup outcome, which are -// the two halves of one story: the event records the refusal cuDNN just handed back, and the -// lookup line is a later query answered from that stored refusal instead of by building the -// graph again. Tell them apart by the line shape -- the event line carries counters, the -// lookup line carries the cache key. +// level 1 (events) : one line per event that happens once per distinct cache key +// (CREATE_GRAPH, BUILD_PLANS), plus the exit summary block and +// its stage timings. Low volume by construction. +// level 2 (trace) : adds a line per cache lookup (HIT/MISS, with the normalized +// key) and per execution (EXEC). High volume, and it serializes +// threads on the stderr lock, which the stage timings are then +// measured under -- no timed region writes to stderr, so they +// stay sound, but they read a little high. // -// An optional ":" suffix picks which processes emit, defaulting to rank 0 -// so that output does not scale with the world size: "1:all" for every rank, -// "2:0,3" for a specific set. See `rank_selected` for when overriding pays off. +// Counters are kept per build site -- f16/fp8 crossed with fwd/bwd -- since one +// process can drive both backends, and every event name is also the counter column +// it increments. What the columns mean, the identities they satisfy and the ratios +// worth reading are with the counter definitions below. // -// Level 1 on one training step of a supported configuration. Every line begins with -// "[FUSED-ATTN-CACHE] rank= | ", or with just "[FUSED-ATTN-CACHE] " when the launcher -// exports no rank (see `rank_tag`), elided below, and carries the running totals, of which -// only the pass being reported is shown (the counters are printed right-aligned in a fixed -// width, and are abbreviated here): +// One level-1 training step, line prefixes and trailing columns elided: // -// f16 fwd BUILD_GRAPH | tid=0 dev=0 | fwd hit_supported=0, miss=1, build_graph=1, ... -// f16 bwd BUILD_GRAPH | tid=0 dev=0 | fwd ... | bwd hit_supported=0, miss=1, ... -// f16 fwd BUILD_PLANS | tid=0 dev=0 | fwd hit_supported=1, miss=1, build_plans=1, ... +// tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, ... // ===== summary begin ===== -// f16 SUMMARY-TID | tid=0 dev=0 | fwd hit_supported=5, miss=1, build_graph=1, ... -// f16 SUMMARY-TID | tid=1 dev=0 | fwd ... | bwd hit_supported=4, build_plans=1, ... -// f16 SUMMARY | tid=all dev=all | fwd hit_supported=5, miss=1, build_graph=1, ... -// f16 fwd check_support | calls=1 | time= 0.031 ms/call +// tid=0 dev=0 | f16 fwd | hit=5, miss=1, create_graph=1, ... +// tid=1 dev=0 | f16 bwd | hit=4, build_plans=1, exec=1, ... +// tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, ... // f16 fwd build_plans | calls=1 | time= 262.104 ms/call // ===== summary end ===== // -// The two thread rows are what a PyTorch step really looks like: the forward, and the support -// probe for the backward, run on the main thread, while the backward itself runs on the -// autograd thread and finds the graph that probe left behind. Neither row satisfies -// `build_graph >= build_plans` by itself -- tid=1 compiled the plans of a graph tid=0 built -- -// so read the identities off the totals rows rather than the per-thread ones. -// -// The device column matters as soon as one process drives more than one -- device_id is part of -// the cache key, so the same shape on two devices is two entries, and a build count that looks -// doubled is explained by reading which device each BUILD_GRAPH came from. -// -// A support query misses and builds, and every later lookup of that key is a hit_supported -- -// including the workspace-sizing call that precedes each execution -- so the hit columns climb -// faster than exec. `build_graph=1, build_plans=1` says that graph went on to be executed; -// `build_graph` above `build_plans` counts graphs built for a query and never run. A refused -// configuration reads `miss=1, unsupported=1, build_graph=0` instead, and stays at one refusal -// however many times it is queried: the repeat queries land in hit_unsupported. -// -// Level 2 adds one line per lookup and per execution, with the key that decided it: -// -// f16 fwd MISS | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... -// f16 fwd HIT | tid=0 dev=0 | train=1 det=0 cg=0 ... b=2 h=16 sq=512 skv=512 ... -// -// where diffing two MISS lines names the fields that cost the extra build. +// Rows for a site a thread never reached are left out rather than zeroed, which is +// why tid=1 has a backward row and no forward one: in a PyTorch step the forward and +// the backward's support probe run on the main thread, and the backward itself on the +// autograd thread, which finds the graph that probe left behind. That split is why +// the build identities hold on the totals rows and not on any single thread's. // ============================================================================ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ @@ -128,15 +92,11 @@ inline int launcher_rank() { return rank; } -// Whether this process emits diagnostics. Every rank writes to the same stderr, -// so emitting from all of them multiplies the volume by the world size -- and -// under data/tensor parallelism the ranks are running identical shapes, so the -// copies say the same thing. Hence rank 0 only by default. -// -// Context parallelism is the case worth overriding for: the ranks run different -// subsets of the per-step regimes (under p2p, rank 0 never sees the lower-triangle -// config that the last rank does), so their build counts genuinely differ. -// Select with the ":" suffix, e.g. "1:all" or "2:0,3". +// Whether this process emits diagnostics. Every rank writes to the same stderr, and under +// data/tensor parallelism they run identical shapes, so emitting from all of them multiplies the +// volume by the world size to say the same thing. Hence rank 0 only by default, overridable with +// the ":" suffix. Context parallelism is the case worth overriding for: the ranks run +// different subsets of the per-step regimes, so their build counts genuinely differ. inline bool rank_selected() { static const bool selected = [] { const int rank = launcher_rank(); @@ -159,13 +119,11 @@ inline bool rank_selected() { return selected; } -// Diagnostics are on at level >= 1, and only for the selected ranks. Unselected -// ranks skip the counters too, so they pay nothing beyond this check. -// -// Cached in its own flag rather than recomputed from the two above, so that this -- the check -// every call site makes, on the per-lookup path included -- reads one initialized-once static -// instead of two. Both inputs are fixed for the life of the process, so there is nothing to -// recompute; `rank_selected` is still only reached when the level says diagnostics are on. +// Diagnostics are on at level >= 1, and only for the selected ranks. Unselected ranks skip the +// counters too, so they pay nothing beyond this check. Cached in its own flag rather than +// recomputed from the two above, so that the check every call site makes -- the per-lookup path +// included -- reads one initialized-once static instead of two. Both inputs are fixed for the +// life of the process. inline bool enabled() { static const bool on = debug_level() >= 1 && rank_selected(); return on; @@ -174,11 +132,10 @@ inline bool enabled() { // Per-lookup / per-exec trace lines are gated behind level >= 2. inline bool trace_enabled() { return debug_level() >= 2; } -// Names the emitting rank. Distributed runs put one process per rank on the same stderr, so -// without this the ranks' lines would be indistinguishable. A run whose launcher exports no -// rank is left untagged rather than falling back to a pid: an OS-level identifier is only -// useful for correlating against a profiler or another process, which these logs are not for. -// The tag carries its own trailing separator, so the untagged case prints no empty column. +// Names the emitting rank, without which the ranks sharing one stderr would be indistinguishable. +// A run whose launcher exports no rank is left untagged rather than falling back to a pid, an +// OS-level identifier only being useful for correlating against a profiler. The tag carries its +// own trailing separator, so the untagged case prints no empty column. inline const std::string &rank_tag() { static const std::string *tag = [] { const int rank = launcher_rank(); @@ -188,10 +145,9 @@ inline const std::string &rank_tag() { return *tag; } -// More readable, shorter thread IDs (0, 1, 2, ...). These are assignment order, not identity: -// tid=0 is whichever thread touched this cache first, and the number means nothing outside this -// process. It exists to attribute the per-thread SUMMARY rows, not to be matched against -// anything external. +// Short thread IDs (0, 1, 2, ...) in assignment order, not identity: tid=0 is whichever thread +// touched this cache first, and the number means nothing outside this process. It attributes the +// per-thread summary rows and is not meant to be matched against anything external. inline unsigned thread_seq_id() { static std::atomic next{0}; static thread_local unsigned id = next.fetch_add(1, std::memory_order_relaxed); @@ -203,19 +159,17 @@ inline unsigned thread_seq_id() { inline void register_summary_once(); // ============================================================================ -// The build site an event came from: f16 or fp8, forward or backward. Every recorder names -// both halves, because the counters are kept per site rather than per pass. One process can -// drive both backends, and adding f16's builds into the same column as fp8's would leave such -// a run unable to say which of them paid for what. +// The build site an event came from: f16 or fp8, forward or backward. Every recorder names both +// halves, since the counters are per site -- adding f16's builds into fp8's column would leave a +// run that drove both unable to say which paid for what. A pair of enums rather than the +// "fwd"/"bwd" strings this used to take also turns a mistake at a call site into a compile error. // -// Backend::F16 is the arbitrary-seqlen f16 backend; the max512 one keeps no graph cache and so -// has nothing to report here. Naming the site with a pair of enums rather than with the -// "fwd"/"bwd" strings this used to take is also what turns a mistake at a call site into a -// compile error instead of an event silently counted against the wrong column. +// Backend::F16 is the arbitrary-seqlen backend; the max512 one keeps no graph cache. Pass is +// fused_attn::Pass, from config_and_params.h, so that a recorder and the key it prints share one +// notion of direction. // ============================================================================ enum class Backend { F16, FP8 }; -enum class Pass { Fwd, Bwd }; inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } @@ -230,85 +184,67 @@ inline constexpr size_t site_index(Backend b, Pass p) { // ============================================================================ // Cache event counters, one block per build site. Each name is both the event tag on the line // that records it and the column carrying its running total: -// - build_graph: a graph built and cached in response to a cache miss. Built only as far -// as check_support(), which is all a support probe needs. -// - build_plans: a cached graph finished with build_plans(), the kernel compilation that -// build_graph deferred. At most one per build_graph, and paid by the first -// execution of that graph rather than by the probe that built it. -// - unsupported: a configuration cuDNN refused, now remembered as a negative cache entry. -// The other way a miss can end. Counted once per refusal recorded, which is -// normally once per distinct refused key; later queries for it are -// hit_unsupported. -// - exec: a graph execution call with valid runtime tensors -// - hit_supported: a lookup answered from the graph map. May not lead to an exec: it can be -// a backend availability check, or the workspace-sizing call of -// nvte_fused_attn_fwd/bwd, which has no runtime tensors to run with. -// - hit_unsupported: a lookup answered from the refusal map -- a key cuDNN has already -// refused, replayed instead of rebuilt. Both hit columns are named for the -// map that answered them, and `unsupported` above counts the refusals -// themselves rather than the queries that replay them. -// - miss: a lookup neither map answered; triggers a graph build +// - create_graph: a graph created and cached for a miss, only as far as check_support(). +// - build_plans: a cached graph finished with graph.build_plans(), the kernel compilation that +// create_graph deferred. At most one per create_graph, paid by that graph's first execution +// rather than by the probe that built it. +// - exec: a graph execution call with valid runtime tensors. +// - hit: a lookup answered from the cache. Need not lead to an exec -- it can be a backend +// availability check, or the workspace-sizing call of nvte_fused_attn_fwd/bwd, which has no +// runtime tensors to run with. +// - miss: a lookup the cache did not answer; triggers a graph build. // -// Identities. These hold by construction, so a violation is a bug in the cache or in the -// counting rather than something the workload did: -// - hit_supported + hit_unsupported + miss = every lookup, one recorded per entry into -// build_or_get_cached_graph, which makes it the denominator for everything below. -// - miss = build_graph + unsupported. A shortfall in either means a build ended in -// something cuDNN did not state as a verdict on the graph. -// - build_graph >= build_plans, the gap being graphs a probe built that nothing has run. -// Eviction grows both rather than closing it: a rebuilt key gets a fresh once_flag. -// - exec > 0 implies build_plans > 0, every site calling ensure_plans_built ahead of the -// workspace-sizing return, which is itself ahead of record_exec. The same ordering read -// backwards: a workspace-sizing call pays build_plans and never exec. -// - hit_unsupported > 0 implies unsupported > 0, a refusal being replayable only once some -// earlier call has recorded it. -// - The two build identities are properties of the totals rows, not of one SUMMARY-TID row: -// the thread that builds a graph need not be the thread that compiles its plans, and a -// PyTorch step splits exactly that way across the autograd thread. -// - A backend's SUMMARY-TID rows sum column by column to its SUMMARY row, and the -// per-backend rows to the all-backends one. -// - A lost build race disturbs none of the above: the loser records its own miss and its own -// build_graph, so both sides of miss = build_graph + unsupported move together, and the -// entry's once_flag still permits only one build_plans. What a race does break is reading -// build_graph as the number of graphs cached, two builds being able to stand behind one -// entry; the same goes for unsupported and the number of keys cuDNN has refused. -// - In the stage timing rows, calls only fall along the sequence validate >= -// build_operation_graph >= create_execution_plans >= check_support, each drop being the -// builds that ended at the stage before -- which localizes where cuDNN refuses, rather -// than only how long refusing took. -// - The build_plans timing row can show more calls than the build_plans column counts, the -// difference being plan builds that threw: the timer records while unwinding, the counter -// only after the call returns. +// Identities, holding by construction, so a violation is a bug in the cache or in the counting +// rather than something the workload did: +// - hit + miss = every lookup, one per entry into lookup_or_cache_graph, which makes it the +// denominator for everything below. +// - miss >= create_graph, the difference being builds that threw, whether cuDNN refused the graph +// or could not reach a verdict. Nothing is cached for those, so this is the only place a +// refusal shows up; its reason goes to the framework instead. +// - create_graph >= build_plans, the gap being graphs a probe built that nothing has run. +// - exec > 0 implies build_plans > 0, every site calling build_plans() ahead of the +// workspace-sizing return, itself ahead of record_exec. Read backwards: a workspace-sizing +// call pays build_plans and never exec. +// - Both build identities belong to the totals rows, not to one thread's: the thread that builds +// a graph need not compile its plans, and a PyTorch step splits exactly that way. +// - Per-thread rows sum column by column to "tid=all dev=all", and the per-backend rows of one +// pass to that pass's all-backends row. +// - A lost build race disturbs none of the above -- the loser records its own miss and its own +// create_graph, and the once_flag still permits one build_plans -- but it does break reading +// create_graph as the number of graphs cached. +// - Stage timing calls fall along validate >= build_operation_graph >= create_execution_plans >= +// check_support, each drop being the builds that ended at the stage before, which localizes +// where cuDNN refuses rather than only how long refusing took. +// - The build_plans timing row can show more calls than the build_plans column, the difference +// being plan builds that threw: the timer records while unwinding, the counter only on return. // -// Signatures. Workload-dependent, so these are read rather than asserted: -// - After warmup only hit_supported and exec should move. A build_graph late in a run means -// something varies per step that need not. -// - Several hit_supported per exec is normal, since backend selection, workspace sizing and -// execution all look the same key up; what matters is that the ratio stays flat. -// - exec / build_graph is the amortization figure, how many executions each built graph -// served, and a lower bound at that, since a race or an eviction adds a build without -// adding a graph. Single digits after a long run means the cache is not earning its keep. -// - hit_unsupported climbing while unsupported stays at one is the negative cache doing its -// job. It also says this site never runs fused, which makes it the column to reach for +// Signatures, workload-dependent, so read rather than asserted: +// - After warmup only hit and exec should move; a late create_graph means something varies per +// step that need not. +// - Several hits per exec is normal, since selection, workspace sizing and execution all look +// the same key up; what matters is that the ratio stays flat. +// - exec / create_graph is the amortization figure, and a lower bound at that, a lost race adding +// a build without a graph. Single digits after a long run means the cache is not earning its +// keep. +// - miss climbing while create_graph stays put is a configuration cuDNN keeps refusing, each query +// paying a discarded build. It also says this site never runs fused, making it the pair to read // when attention is slower than expected and nothing raised an error. -// - build_graph or unsupported past kCacheCapacity suggests that map has evicted. A hint -// rather than an identity: a lost build race counts twice against one key. -// - Two MISS lines carrying the same key, with build_graph above the number of distinct keys, -// is that lost race. It is wasted work rather than a bug, and worth chasing only if it -// repeats, which would mean threads are arriving on cold keys together every step. -// - A level-2 trace is the set of lookups that happened, not the order they happened in: the -// line is written after the cache lock is dropped, so two threads that raced for it can -// print in the opposite order. +// - miss climbing without settling means the key space is not closing, and since the cache is +// unbounded, every distinct key is held for the life of the process. +// - A build count that looks doubled on a multi-device process usually is not: device_id is part +// of the key, so the same shape on two devices is two entries. Read the dev column. +// - Two MISS lines with the same key, create_graph above the number of distinct keys, is that lost +// race: wasted work rather than a bug, worth chasing only if it repeats. +// - A level-2 trace is the set of lookups, not their order, the line being written after the +// cache lock is dropped. // ============================================================================ struct EventCounters { - std::atomic build_graph{0}; + std::atomic create_graph{0}; std::atomic build_plans{0}; std::atomic exec{0}; - std::atomic hit_supported{0}; - std::atomic hit_unsupported{0}; + std::atomic hit{0}; std::atomic miss{0}; - std::atomic unsupported{0}; }; inline EventCounters &counters(Backend b, Pass p) { @@ -316,75 +252,59 @@ inline EventCounters &counters(Backend b, Pass p) { return table[site_index(b, p)]; } -// One counter block read out into plain values. The summary sums blocks to get its per-backend -// and all-backends rows, atomics cannot be summed, and this is where the reading happens; it -// also keeps the loads out of the formatting. The columns are not read as one indivisible -// operation, which nothing here wants: the summary runs at exit, after the threads that wrote -// them are done, and an event line is a snapshot of a moving count by nature. +// One counter block read out into plain values, so the summary can sum blocks for its per-backend +// and all-backends rows. The columns are not read as one indivisible operation, which nothing here +// wants: the summary runs at exit, after the writing threads are done, and an event line is a +// snapshot of a moving count by nature. struct CounterSnapshot { - uint64_t build_graph = 0; + uint64_t create_graph = 0; uint64_t build_plans = 0; uint64_t exec = 0; - uint64_t hit_supported = 0; - uint64_t hit_unsupported = 0; + uint64_t hit = 0; uint64_t miss = 0; - uint64_t unsupported = 0; CounterSnapshot &operator+=(const CounterSnapshot &other) { - build_graph += other.build_graph; + create_graph += other.create_graph; build_plans += other.build_plans; exec += other.exec; - hit_supported += other.hit_supported; - hit_unsupported += other.hit_unsupported; + hit += other.hit; miss += other.miss; - unsupported += other.unsupported; return *this; } // Whether this block saw nothing at all, which is what lets the summary leave out the rows // for a backend the run never used rather than printing zeros for it. - bool empty() const { - return (build_graph | build_plans | exec | hit_supported | hit_unsupported | miss | - unsupported) == 0; - } + bool empty() const { return (create_graph | build_plans | exec | hit | miss) == 0; } }; inline CounterSnapshot snapshot(const EventCounters &c) { CounterSnapshot s; - s.build_graph = c.build_graph.load(std::memory_order_relaxed); + s.create_graph = c.create_graph.load(std::memory_order_relaxed); s.build_plans = c.build_plans.load(std::memory_order_relaxed); s.exec = c.exec.load(std::memory_order_relaxed); - s.hit_supported = c.hit_supported.load(std::memory_order_relaxed); - s.hit_unsupported = c.hit_unsupported.load(std::memory_order_relaxed); + s.hit = c.hit.load(std::memory_order_relaxed); s.miss = c.miss.load(std::memory_order_relaxed); - s.unsupported = c.unsupported.load(std::memory_order_relaxed); return s; } -// Per-thread counters, one block per build site, so the summary can break down every column by -// thread and backend. In the single-process context-parallel case each device is driven by its -// own thread, so this reveals which thread built and executed what; under PyTorch it also -// separates the main thread from the autograd thread that runs the backward. +// Per-thread counters, one block per build site, so the summary can break every column down by +// thread and backend: in the single-process context-parallel case each device is driven by its own +// thread, and under PyTorch this separates the main thread from the autograd one. // -// `device` is the device this thread last drove, restamped on every event. The event lines print -// the live current device, which is exact; this exists for the SUMMARY-TID rows, which are -// written at exit by whichever thread is exiting and so cannot ask the recorded thread what it -// was working on. A thread that stays on one device -- which is the arrangement everything here -// is built around, device_id being part of the cache key -- makes the two the same answer. +// `device` is the device this thread last drove, restamped on every event. Event lines print the +// live current device instead, which is exact; this exists for the per-thread summary rows, written +// at exit by whichever thread is exiting, which cannot ask the recorded thread what it was doing. struct ThreadCounters { unsigned tid = 0; std::atomic device{-1}; std::array sites; }; -// The registry and its mutex are heap-allocated and deliberately never freed. -// Function-local static destructors and atexit handlers run as a single sequence, -// in reverse order of construction/registration. This registry is built lazily, so -// it can be constructed *after* the summary handler is registered -- in which case -// it would be destroyed *before* that handler runs, leaving the handler to lock a -// destroyed mutex and walk a destroyed vector. Leaking removes the ordering -// question rather than reasoning about it, and the cost is bounded: one mutex and -// one vector for the process, reclaimed by the OS at exit anyway. +// The registry and its mutex are heap-allocated and deliberately never freed. Static destructors +// and atexit handlers run as one sequence in reverse order of construction, and this registry is +// built lazily, so it can be constructed *after* the summary handler is registered -- and would +// then be destroyed *before* it runs, leaving the handler to lock a destroyed mutex and walk a +// destroyed vector. Leaking removes the ordering question, at a cost of one mutex and one vector. inline std::mutex &thread_registry_mutex() { static std::mutex *m = new std::mutex(); return *m; @@ -394,17 +314,15 @@ inline std::vector &thread_registry() { return *v; } -// This thread's counter block, leaked for a related but distinct reason: a worker -// thread can exit long before the process does, while the registry keeps a pointer -// to its block for the end-of-run summary. Tying the block's lifetime to the -// thread would leave that pointer dangling. One small struct per thread. +// This thread's counter block, leaked for a related but distinct reason: a worker thread can exit +// long before the process does, while the registry holds a pointer to its block for the exit +// summary. Tying the block's lifetime to the thread would leave that pointer dangling. inline ThreadCounters &thread_counters() { static thread_local ThreadCounters *tc = [] { auto *p = new ThreadCounters(); p->tid = thread_seq_id(); - // Stamped here as well as on every event, so that a thread which only ever hits the cache -- - // and so never reaches print_counters() at level 1 -- still names a device in the summary - // rather than reporting the -1 it was constructed with. + // Stamped here as well as on every event, so a thread that only ever hits the cache -- never + // reaching print_counters() at level 1 -- still names a device rather than the -1 it began at. p->device.store(cuda::current_device(), std::memory_order_relaxed); { std::lock_guard lock(thread_registry_mutex()); @@ -419,71 +337,69 @@ inline EventCounters &thread_counters(Backend b, Pass p) { return thread_counters().sites[site_index(b, p)]; } -// Format one pair of counter blocks -- the two passes of a single backend -- as one line. -// `label` is the event or summary tag, and names the backend whenever the line speaks for one. -// `tid_field` is the whole thread column, e.g. "tid=3"; the totals rows pass "tid=all" so that -// they cannot be misread as thread 0's row. `dev_field` is the device column and works the same -// way, "dev=all" on a totals row -- those counters are summed across whatever devices the -// process drove, so naming one of them would be a lie. +// Format one counter block -- one pass of one backend -- as one line. One pass rather than both +// because a line carrying the forward and backward columns together ran past 300 characters and +// wrapped in most terminals; the two passes are adjacent rows instead. +// +// `tid_field` and `dev_field` are whole columns, e.g. "tid=3" and "dev=0". The totals rows pass +// "tid=all" and "dev=all", since those counters are summed across whatever the process drove and +// naming one thread or device would be a lie. // -// What the columns mean, the identities they can be asserted against and the ratios worth -// reading are all with the counter definitions above. -inline std::string format_counter_line(const char *label, const char *tid_field, - const char *dev_field, const CounterSnapshot &f, - const CounterSnapshot &b) { - char buf[768]; +// `label` is the build site, "f16 fwd", plus the event name on an event line, and arrives padded to +// the width its own kind of line uses: 20 characters for an event line, 7 for a summary row. +// Deliberately not one width for both -- sharing it would put twelve blank columns on every summary +// row to align the scattered event lines against a block that is delimited and read on its own. +// +// The thread and device come first, so every line, level-2 trace lines included, shares one prefix +// to read down. What the columns mean and the identities they satisfy are with the definitions +// above. +inline std::string format_counter_line(const char *tid_field, const char *dev_field, + const char *label, const CounterSnapshot &c) { + char buf[512]; std::snprintf(buf, sizeof(buf), - "[FUSED-ATTN-CACHE] %s%-19s | %-7s %-7s | fwd hit_supported=%4" PRIu64 - ", hit_unsupported=%4" PRIu64 ", miss=%4" PRIu64 ", build_graph=%4" PRIu64 - ", unsupported=%4" PRIu64 ", build_plans=%4" PRIu64 ", exec=%4" PRIu64 - " | bwd hit_supported=%4" PRIu64 ", hit_unsupported=%4" PRIu64 ", miss=%4" PRIu64 - ", build_graph=%4" PRIu64 ", unsupported=%4" PRIu64 ", build_plans=%4" PRIu64 - ", exec=%4" PRIu64 "\n", - rank_tag().c_str(), label, tid_field, dev_field, f.hit_supported, f.hit_unsupported, - f.miss, f.build_graph, f.unsupported, f.build_plans, f.exec, b.hit_supported, - b.hit_unsupported, b.miss, b.build_graph, b.unsupported, b.build_plans, b.exec); + "[FUSED-ATTN-CACHE] %s%-7s %-7s | %s | hit=%4" PRIu64 ", miss=%4" PRIu64 + ", create_graph=%4" PRIu64 ", build_plans=%4" PRIu64 ", exec=%4" PRIu64 "\n", + rank_tag().c_str(), tid_field, dev_field, label, c.hit, c.miss, c.create_graph, + c.build_plans, c.exec); return std::string(buf); } -inline void print_counter_block(const char *label, const char *tid_field, const char *dev_field, - const CounterSnapshot &f, const CounterSnapshot &b) { - const std::string line = format_counter_line(label, tid_field, dev_field, f, b); - std::fputs(line.c_str(), stderr); - std::fflush(stderr); -} - -// One event line, from the thread the event happened on, carrying the running totals of the -// backend that raised it. The device is read live rather than remembered, so it is the device -// this event was actually issued against, and is recorded on the thread's block on the way past -// for the benefit of the exit summary. +// One event line, from the thread the event happened on, carrying the running totals of the build +// site that raised it. The device is read live rather than remembered, so it is the device this +// event was actually issued against, and is recorded on the thread's block on the way past for +// the benefit of the exit summary. inline void print_counters(Backend b, Pass p, const char *event) { const int device = cuda::current_device(); thread_counters().device.store(device, std::memory_order_relaxed); char label[32]; char tid_field[16]; char dev_field[16]; - std::snprintf(label, sizeof(label), "%s %s %s", backend_name(b), pass_name(p), event); + // The event name is padded to the longest of them, so that the counters of one event line fall + // where the next one's do. + std::snprintf(label, sizeof(label), "%s %s %-12s", backend_name(b), pass_name(p), event); std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); - print_counter_block(label, tid_field, dev_field, snapshot(counters(b, Pass::Fwd)), - snapshot(counters(b, Pass::Bwd))); + const std::string line = + format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p))); + std::fputs(line.c_str(), stderr); + std::fflush(stderr); } -// A graph built through check_support() and cached. Call after the build, from the miss -// path that performed it. -inline void record_graph_built(Backend b, Pass p) { +// A graph created, taken through check_support() and cached. Call after that, from the miss path +// that did it -- after, because a graph cuDNN refuses throws instead of arriving here, which is +// what makes miss - create_graph the count of refused builds. +inline void record_graph_created(Backend b, Pass p) { if (!enabled()) return; register_summary_once(); - counters(b, p).build_graph.fetch_add(1, std::memory_order_relaxed); - thread_counters(b, p).build_graph.fetch_add(1, std::memory_order_relaxed); - print_counters(b, p, "BUILD_GRAPH"); + counters(b, p).create_graph.fetch_add(1, std::memory_order_relaxed); + thread_counters(b, p).create_graph.fetch_add(1, std::memory_order_relaxed); + print_counters(b, p, "CREATE_GRAPH"); } -// The build_plans() a build_graph deferred, now completed. Call from inside the std::call_once -// that runs it, after the call returns rather than before: build_plans() throws without -// setting the once_flag, leaving a later execution to retry it, so counting on the way out -// keeps this a count of graphs that reached a runnable state. Like build_graph this fires once -// per distinct cache key, so it stays on the level-1 path. +// The graph.build_plans() a create_graph deferred, now completed. Call from inside the +// std::call_once that runs it, and after the call returns rather than before: it throws without +// setting the once_flag, leaving a later execution to retry, so counting on the way out keeps this +// a count of graphs that reached a runnable state. inline void record_plans_built(Backend b, Pass p) { if (!enabled()) return; register_summary_once(); @@ -492,18 +408,6 @@ inline void record_plans_built(Backend b, Pass p) { print_counters(b, p, "BUILD_PLANS"); } -// A build that cuDNN refused, now remembered as a negative cache entry. Call from the miss path -// that attempted it, in place of record_graph_built(): a refusal and a build are the two ways a -// miss can end, and counting both keeps `miss = build_graph + unsupported` true. Fires once per -// refused key -- later queries for it land in hit_unsupported -- so it stays on the level-1 path. -inline void record_unsupported(Backend b, Pass p) { - if (!enabled()) return; - register_summary_once(); - counters(b, p).unsupported.fetch_add(1, std::memory_order_relaxed); - thread_counters(b, p).unsupported.fetch_add(1, std::memory_order_relaxed); - print_counters(b, p, "UNSUPPORTED"); -} - inline void record_exec(Backend b, Pass p) { if (!enabled()) return; register_summary_once(); @@ -514,19 +418,15 @@ inline void record_exec(Backend b, Pass p) { print_counters(b, p, "EXEC"); } -// What a lookup found. Unsupported is the negative-cache case: a key whose graph cuDNN has -// already refused, so the answer is a remembered refusal rather than a graph. -enum class LookupResult { Miss, Hit, Unsupported }; +// What a lookup found: an entry, or nothing. +enum class LookupResult { Miss, Hit }; -// The column a lookup lands in, which is the cache map that answered it. Written as a switch -// with no default so that adding an outcome fails to compile here rather than being silently -// counted as a miss. +// The column a lookup lands in. Written as a switch with no default so that adding an outcome +// fails to compile here rather than being silently counted as a miss. inline std::atomic &lookup_column(EventCounters &c, LookupResult result) { switch (result) { case LookupResult::Hit: - return c.hit_supported; - case LookupResult::Unsupported: - return c.hit_unsupported; + return c.hit; case LookupResult::Miss: break; } @@ -537,36 +437,24 @@ inline const char *lookup_name(LookupResult result) { switch (result) { case LookupResult::Hit: return "HIT"; - case LookupResult::Unsupported: - return "UNSUPPORTED"; case LookupResult::Miss: break; } return "MISS"; } -// `key` is the normalized cache key -- make_cache_key()'s output, the exact value the -// lookup was performed with -- not the execution config it was derived from. That is -// deliberate: HIT/MISS is decided by comparing keys, so a trace of anything else cannot -// explain its own outcome. Logging the pre-normalization config would show pairs of -// identical lines with opposite outcomes (normalization having collapsed a difference, -// e.g. bottom_right_diagonal or the THD token counts) and pairs of differing lines that -// both hit (the difference being in a field the key drops, e.g. attn_scale). Diffing two -// MISS lines here instead names exactly the fields responsible for the extra build. +// `key` is the normalized cache key -- make_cache_key(pass)'s output, the exact value looked up -- +// not the execution config it came from. HIT/MISS is decided by comparing keys, so a trace of +// anything else cannot explain its own outcome: the pre-normalization config would show identical +// lines with opposite outcomes, and differing lines that both hit. Diffing two MISS lines here +// names exactly the fields responsible for the extra build. // -// The cost is that fields normalization overwrites are no longer visible in their -// original form: attn_scale reads 1, ragged num_tokens read 0, and max_seqlen/batch_size -// read their bucketed values. Recover those from the caller if a line needs to be traced -// back to a specific test case. +// The cost is that overwritten fields are no longer visible in their original form: attn_scale +// reads 1, ragged num_tokens read 0, max_seqlen and batch_size read their bucketed values. inline void record_cache_lookup(Backend b, Pass p, LookupResult result, const FusedAttnConfig &key) { if (!enabled()) return; register_summary_once(); - // A refusal replayed from the negative cache is counted apart from a graph hit, in - // hit_unsupported rather than in hit_supported. Both were answered without building - // anything, which is what the two hit columns have in common; which map answered is the - // thing worth being able to read off a level-1 summary, since a run whose hits are mostly - // replayed refusals is not reusing graphs at all. lookup_column(counters(b, p), result).fetch_add(1, std::memory_order_relaxed); lookup_column(thread_counters(b, p), result).fetch_add(1, std::memory_order_relaxed); // The per-lookup config dump is the highest-volume line (one per cache lookup); @@ -574,7 +462,7 @@ inline void record_cache_lookup(Backend b, Pass p, LookupResult result, if (!trace_enabled()) return; std::fprintf( stderr, - "[FUSED-ATTN-CACHE] %s%-3s %-3s %-11s | tid=%u dev=%d | train=%d det=%d cg=%d " + "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | 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 @@ -584,8 +472,8 @@ inline void record_cache_lookup(Backend b, Pass p, LookupResult result, " 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", - rank_tag().c_str(), backend_name(b), pass_name(p), lookup_name(result), thread_seq_id(), - key.device_id, static_cast(key.is_training), static_cast(key.deterministic), + rank_tag().c_str(), thread_seq_id(), key.device_id, backend_name(b), pass_name(p), + lookup_name(result), static_cast(key.is_training), static_cast(key.deterministic), static_cast(key.cuda_graph), static_cast(key.return_max_logit), static_cast(key.check_for_forward_support), static_cast(key.attn_mask_type), static_cast(key.bias_type), static_cast(key.window_size_left), @@ -614,21 +502,17 @@ inline void record_cache_lookup(Backend b, Pass p, LookupResult result, // ============================================================================ // Graph build timings. // -// A cuDNN graph build is a fixed sequence of frontend calls, and which one -// dominates determines what to do about a slow build: time in `check_support` -// and `build_plans` is heuristic selection and kernel compilation, largely -// intrinsic to the shape, whereas time in `validate` or `build_operation_graph` -// is graph-construction cost on our side of the boundary. Timing the stages -// separately is what makes that distinction; one duration per build cannot. +// Which stage dominates determines what to do about a slow build: time in +// `check_support` and `build_plans` is heuristic selection and kernel compilation, +// largely intrinsic to the shape, while time in `validate` or +// `build_operation_graph` is graph-construction cost on our side. One duration per +// build cannot make that distinction. // -// Each stage is wrapped where it is called -- graph_cache.h, which is where all -// five frontend calls live -- and accumulates into the table below, under the -// pass its caller was serving. -// The end-of-run summary reports each as a mean over its calls. Only sums are -// kept, so the mean is all that can be recovered -- and since a build happens -// once per distinct cache key, those calls span different shapes rather than -// repeating one. Read a stage mean as where build time goes in aggregate, not as -// the cost of any particular build. +// Each stage is wrapped where it is called, in graph_cache.h, and accumulates into +// the table below under its build site. Only sums are kept, so the summary can +// report a mean and nothing else -- and since a build happens once per distinct +// cache key, those calls span different shapes rather than repeating one. Read a +// stage mean as where build time goes in aggregate, not as any one build's cost. // ============================================================================ // The frontend calls that make up a build, in the order they run. `kCount` must @@ -661,13 +545,10 @@ inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { // Times one stage: clock read in the constructor, accumulated in the destructor. // Recording on scope exit rather than at an explicit stop() keeps a failing stage -// measurable: `build_plans` throws through NVTE_CHECK_CUDNN_FE and the destructor -// still runs during unwinding, so a build that dies there contributes its time to -// the failure instead of vanishing from the summary. The four stages before it -// return their status instead of throwing, and are timed the same way for the same -// reason. `on` is latched at construction rather than re-tested in the destructor, -// which is what keeps that symmetric: the destructor can never accumulate against a -// `start` the constructor left unset. +// measurable, since `build_plans` throws through NVTE_CHECK_CUDNN_FE and the +// destructor still runs while unwinding, so a build that dies there contributes its +// time instead of vanishing. `on` is latched at construction rather than re-tested in +// the destructor, so the destructor can never accumulate against an unset `start`. struct ScopedBuildTimer { BuildStage stage; bool on; @@ -691,12 +572,12 @@ struct ScopedBuildTimer { } }; -// Time `fn` as `stage` of the given build site, named as the record_* helpers above name it. -// Preferred over declaring a ScopedBuildTimer at the call site: the measured region is exactly -// the call passed in, so surrounding work cannot drift into it as that code changes. With -// diagnostics off this costs one cached-flag check, and that is per build rather than per lookup. +// Record how long `fn` takes as `stage` of the given build site. Unlike the record_* helpers above +// this wraps the work rather than reporting on work already done, which is the point: preferred +// over a ScopedBuildTimer at the call site because the measured region is exactly the call passed +// in, so surrounding work cannot drift into it as that code changes. template -inline void timer(Backend b, Pass p, BuildStage stage, Fn &&fn) { +inline void record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { ScopedBuildTimer scoped(b, p, stage); fn(); } @@ -708,14 +589,12 @@ inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; - // Build the whole summary in memory and emit it with a single write, so - // that the blocks of concurrently-exiting processes (one per rank under - // torchrun) stay grouped instead of interleaving line by line. + // Built in memory and emitted with one write, so that concurrently-exiting + // processes (one per rank under torchrun) stay grouped rather than interleaving. std::string block; block += "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary begin =====\n"; constexpr Backend kBackends[] = {Backend::F16, Backend::FP8}; - // A backend the run never reached is left out of the summary rather than reported as a - // row of zeros, so the usual single-backend run reads as it did before this was split. + // A backend the run never reached is left out rather than reported as a row of zeros. size_t active_backends = 0; for (const Backend b : kBackends) { if (!snapshot(counters(b, Pass::Fwd)).empty() || @@ -723,8 +602,8 @@ inline void register_summary_once() { ++active_backends; } } - // Per-thread breakdown (sorted by tid), one row per backend that thread drove. Useful in - // the single-process context-parallel case where each device runs on its own thread. + // Per-thread breakdown (sorted by tid), one row per build site that thread drove, with + // unreached sites left out for the same reason an unused backend is. { std::lock_guard lock(thread_registry_mutex()); std::vector blocks = thread_registry(); @@ -737,32 +616,40 @@ inline void register_summary_once() { std::snprintf(dev_field, sizeof(dev_field), "dev=%d", tc->device.load(std::memory_order_relaxed)); for (const Backend b : kBackends) { - const CounterSnapshot fwd = snapshot(tc->sites[site_index(b, Pass::Fwd)]); - const CounterSnapshot bwd = snapshot(tc->sites[site_index(b, Pass::Bwd)]); - if (fwd.empty() && bwd.empty()) continue; - char label[32]; - std::snprintf(label, sizeof(label), "%s SUMMARY-TID", backend_name(b)); - block += format_counter_line(label, tid_field, dev_field, fwd, bwd); + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot c = snapshot(tc->sites[site_index(b, p)]); + if (c.empty()) continue; + // No padding: a site name is exactly the width of the column on a summary row. + char label[32]; + std::snprintf(label, sizeof(label), "%s %s", backend_name(b), pass_name(p)); + block += format_counter_line(tid_field, dev_field, label, c); + } } } } - // Totals last, so they read as the sum of the per-thread rows above: one row per backend, - // then a row across the backends only when the run used more than one. With a single - // backend that row would repeat the one above it verbatim and say nothing extra. + // Totals last, so they read as the sum of the per-thread rows above: one row per build site, + // then a row per pass across the backends only when the run used more than one, since with a + // single backend those would repeat the rows above verbatim. CounterSnapshot all_fwd; CounterSnapshot all_bwd; for (const Backend b : kBackends) { - const CounterSnapshot fwd = snapshot(counters(b, Pass::Fwd)); - const CounterSnapshot bwd = snapshot(counters(b, Pass::Bwd)); - all_fwd += fwd; - all_bwd += bwd; - if (fwd.empty() && bwd.empty()) continue; - char label[32]; - std::snprintf(label, sizeof(label), "%s SUMMARY", backend_name(b)); - block += format_counter_line(label, "tid=all", "dev=all", fwd, bwd); + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot c = snapshot(counters(b, p)); + (p == Pass::Fwd ? all_fwd : all_bwd) += c; + if (c.empty()) continue; + char label[32]; + std::snprintf(label, sizeof(label), "%s %s", backend_name(b), pass_name(p)); + block += format_counter_line("tid=all", "dev=all", label, c); + } } if (active_backends > 1) { - block += format_counter_line("SUMMARY", "tid=all", "dev=all", all_fwd, all_bwd); + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot &c = (p == Pass::Fwd ? all_fwd : all_bwd); + if (c.empty()) continue; + char label[32]; + std::snprintf(label, sizeof(label), "all %s", pass_name(p)); + block += format_counter_line("tid=all", "dev=all", label, c); + } } for (const Backend b : kBackends) { for (const Pass p : {Pass::Fwd, Pass::Bwd}) { From a82b6bd32b279a716827c9be51a205137d6d557a Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 20 Aug 2026 08:18:34 -0700 Subject: [PATCH 83/88] WIP: tidy up structure Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 6 +- .../attention/test_attention_with_cp.py | 1 - tests/pytorch/utils.py | 11 +- .../common/fused_attn/config_and_params.cpp | 131 +++- .../common/fused_attn/config_and_params.h | 105 ++- .../common/fused_attn/fused_attn.cpp | 111 ++- .../fused_attn_f16_arbitrary_seqlen.cu | 171 ++--- .../fused_attn_f16_arbitrary_seqlen.h | 17 +- .../common/fused_attn/fused_attn_fp8.cu | 274 +++----- .../common/fused_attn/fused_attn_fp8.h | 12 +- .../common/fused_attn/graph_cache.h | 201 ++++-- .../common/fused_attn/graph_cache_debug.h | 657 +++++++++++------- 12 files changed, 932 insertions(+), 765 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index f17842c5e1..85881f512c 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -202,13 +202,13 @@ backend-selection overview. :Default: ``0`` :Description: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). Output goes to stderr, prefixed ``[FUSED-ATTN-CACHE]``. - ``1`` emits one line per event that happens once per distinct cache key -- ``CREATE_GRAPH`` when a graph is constructed, ``BUILD_PLANS`` when its kernels are compiled on first execution -- plus an end-of-run summary block (one row per build site, per thread and in total, plus a row per pass across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries. + ``1`` emits one line per event that happens once per distinct cache key -- ``CREATE_GRAPH`` when a graph is constructed, ``CACHE_GRAPH`` when cuDNN has agreed to run it, ``BUILD_PLANS`` when its kernels are compiled on first execution -- plus an end-of-run summary block (one row per build site, per thread and in total, plus a row per pass across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries. Every line names the thread and device it came from, then the build site behind it -- ``f16`` or ``fp8``, then the pass -- and carries the counters of that site alone, so a process that uses both backends can still tell which of them built what. One line is one pass; the forward and the backward read as adjacent rows. A summary row is an event line without the event name, the block is delimited by ``===== summary begin =====`` and ``===== summary end =====``, and ``tid=all dev=all`` marks the totals. A build site the run never reached is left out rather than shown as a row of zeros. - ``hit`` and ``miss`` account for every lookup, and a miss ends either in ``create_graph`` or in a build cuDNN refused, so ``miss`` minus ``create_graph`` is the number of refusals. Nothing is cached for a configuration cuDNN refuses, so that difference counts refused builds rather than refused configurations: a configuration that is queried again is built and refused again. ``miss`` climbing while ``create_graph`` stays put says this site never runs fused and keeps paying a discarded graph build to find that out, which makes it the pair to read when attention is slower than expected and nothing raised an error. The reason cuDNN gave is not logged here; it reaches the framework as the message explaining why the fused backend was not selected. + ``hit`` and ``miss`` account for every lookup. A miss builds a graph, counted in ``create_graph``, and keeps it only if cuDNN agrees to run it, counted in ``cache_graph``; so the columns fall ``miss`` >= ``create_graph`` >= ``cache_graph``. ``create_graph`` minus ``cache_graph`` is graphs cuDNN refused to run. Nothing is cached for them, so this counts rejected builds rather than rejected configurations: a configuration that is queried again is built and rejected again. A site whose ``cache_graph`` stays put while ``miss`` climbs never runs fused and keeps paying to find that out, which makes these the columns to read when attention is slower than expected and nothing raised an error -- and at level 1, a ``CREATE_GRAPH`` line with no ``CACHE_GRAPH`` after it is one such rejection as it happens. The reason cuDNN gave is not logged here; it reaches the framework as the message explaining why the fused backend was not selected. ``miss`` and ``create_graph`` should agree: configurations FusedAttention itself does not serve are refused before any graph is built, and reported as that same message, so a gap between those two columns means a graph build failed where none was expected to. - ``2`` additionally emits a per-lookup ``HIT``/``MISS`` line carrying the full cache key, and a per-execution ``EXEC`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. + ``2`` additionally emits a per-lookup ``HIT``/``MISS`` line carrying the full cache key, and a per-execution ``EXECUTE`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index ac08436460..6f2efaf83a 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -377,7 +377,6 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type qkv_layout="_".join([qkv_format] * 3), cp_size=num_gpus, cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, - skip_fused_attn=True, ) flash_attn_supported, *_ = available_backends if not flash_attn_supported: diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 90fbcc16b5..cdab36b2c8 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -351,18 +351,11 @@ def get_available_attention_backends( score_mod_bprop: bool = False, cp_size: int = 1, cp_size_a2a: int = 1, - skip_fused_attn: bool = False, ) -> Tuple[List, List]: - """Check for all available attention backends that support a model configuration - - Set `skip_fused_attn=True` to leave fused attention out of the query. The reported - fused-attention backends are then empty, while the FlashAttention and unfused results - are unaffected. This skips cuDNN's support checks, which build and cache a graph per - configuration. - """ + """Check for all available attention backends that support a model configuration""" os.environ["NVTE_FLASH_ATTN"] = "1" - os.environ["NVTE_FUSED_ATTN"] = "0" if skip_fused_attn else "1" + os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True alibi_slopes_shape = None diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index e6ee965b02..8a28ba93ae 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -9,7 +9,10 @@ #include #include +#include +#include #include +#include #include "../common.h" #include "../util/cuda_runtime.h" @@ -54,6 +57,11 @@ void FusedAttnConfig::derive() { 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_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); + is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING); + is_dropout = is_training && dropout != 0.0f; // Both layouts describe variable-length sequences inside padded dimensions, so the mask is the // only thing that tells cuDNN where the real tokens end; without it the graph attends to @@ -76,12 +84,25 @@ void FusedAttnConfig::derive() { 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 + // Use of cu_seqlens vs actual_seqlens, once per backend. Newer cuDNN SDPA can take sequence + // lengths directly as a cumulative tensor, which saves one kernel call; the frontend gates that + // on min(compile-time, runtime) cuDNN, so both versions are tested. The FP8 path needs newer + // versions of both than the F16 path, which is the whole reason there are two answers here. 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; + // Frontend 1.26 supports fp8+cu_seqlens for the C++ API; the Python API needs 1.27. The dropout + // exclusion is not the F16 one restated: the frontend cannot combine dropout with stats + // generation in the fprop unified engine, so such a request would be routed to the old composite + // SDPA engine, which has no cu_seqlens support at all. Remove that term when it can be. + fp8_uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12600 && + (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && + !is_dropout; + + // What each pass stores, classified for the FP8 backend; see the fields for what reads them. + o_is_fp8 = (o_dtype == kNVTEFloat8E4M3 || o_dtype == kNVTEFloat8E5M2); + dqkv_is_fp8 = (dqkv_dtype == kNVTEFloat8E4M3 || dqkv_dtype == kNVTEFloat8E5M2); // packed vs dense dimensions for a ragged (THD) graph; SM8x and SM120 require dense, // BHSD-like dimensions for the Stats/LSE auxiliary tensors and so take the dense path @@ -95,6 +116,28 @@ void FusedAttnConfig::derive() { graph_max_seqlen_kv = (is_ragged_kv && uses_packed_ragged_graph) ? bucketed_num_tokens_kv : max_seqlen_kv; + // Batch size and ragged-offset width the graph is built at, for each direction. One condition + // decides all four: whether cuDNN is handed the caller's cu_seqlens* buffers untouched, which + // only the forward graph ever is. When it is, those buffers are what the graph has to match -- + // their [batch_size + 1] length, which a bucketed batch would read past the end of, and their + // int32 width. The backward graph always reads seqlens converted into our own workspace, so + // neither constraint reaches it and its two answers are the unconditional ones. Otherwise the + // batch is the bucketed one wherever a ragged layout is packed, so that one graph serves every + // batch in its bucket -- the same reason graph_max_seqlen_* stands in for the sequence lengths -- + // and the offset width is whichever the runtime supports, which is what lets older cuDNN + // runtimes work rather than fail. + // + // Kept as four adjacent assignments off shared locals because the pairs have to agree: a forward + // graph built at a bucketed batch while expecting offsets at the other width is the failure this + // arrangement exists to make visible. + const DType wide_ragged_offsets = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; + const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_packed_ragged_graph; + ragged_offset_type_fwd = uses_cu_seqlens_directly ? DType::kInt32 : wide_ragged_offsets; + ragged_offset_type_bwd = wide_ragged_offsets; + graph_batch_size_fwd = + (buckets_the_batch && !uses_cu_seqlens_directly) ? bucketed_batch_size : batch_size; + graph_batch_size_bwd = buckets_the_batch ? bucketed_batch_size : batch_size; + // elements per token for each ragged tensor ragged_offset_mults = RaggedOffsetMultipliers( layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), @@ -125,41 +168,11 @@ void FusedAttnConfig::derive() { is_derived = true; } -GraphDims graph_dims(const FusedAttnConfig &cfg, Pass pass) { - check_derived(cfg); - GraphDims dims; - - // The one condition both answers turn on: the forward graph can be handed the user's cu_seqlens* - // buffers untouched, and then it is those buffers the graph has to match -- their - // [batch_size + 1] length, which a bucketed batch would read past the end of, and their int32 - // width. The backward graph always reads seqlens converted into our own workspace, so nothing - // there is sized by the true batch and nothing there is held to int32. - const bool cudnn_reads_users_cu_seqlens = pass == Pass::Fwd && cfg.uses_cu_seqlens_directly; - - if (cudnn_reads_users_cu_seqlens) { - dims.ragged_offset_type = DType::kInt32; - } else { - // Choose between 32-bit and 64-bit offsets by what the runtime supports, which is what lets - // older cuDNN runtimes work rather than fail. - dims.ragged_offset_type = cudnnGetVersion() >= 90500 ? DType::kInt64 : DType::kInt32; - } - - // Build at the bucketed batch where a ragged layout is packed, so that one graph serves every - // batch in its bucket -- the same reason graph_max_seqlen_* stands in for the sequence lengths. - dims.batch_size = static_cast(cfg.batch_size); - if ((cfg.is_ragged_q || cfg.is_ragged_kv) && cfg.uses_packed_ragged_graph && - !cudnn_reads_users_cu_seqlens) { - dims.batch_size = static_cast(cfg.bucketed_batch_size); - } - - return dims; -} - FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { // Requires a derived config: every normalization below reads a derived field -- is_padding and // is_causal_bottom_right, the is_ragged_* pair, the graph_max_seqlen_* dimensions, and the - // uses_* flags. A precondition rather than an assert, since every caller reaches this through a - // cache_graph_* wrapper, which asserts it once for both the key and the graph. + // uses_* flags. A precondition rather than an assert, since every caller reaches this through + // get_graph(), which asserts it once for both the key and the graph. FusedAttnConfig cache_cfg = *this; // Key the device ID for multi-GPU single-process runs @@ -183,13 +196,14 @@ FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { cache_cfg.max_seqlen_kv = cache_cfg.graph_max_seqlen_kv; // Name the batch size the graph is built at, and drop the token counts the bucketing replaced. - // Asking graph_dims() rather than restating its rule is what keeps the key from naming a batch - // the graph was not built with -- the two directions bucket differently, and the graph builders - // ask the same question with the same pass. + // Reading the batch derive() recorded for this pass, rather than restating the rule that set it, + // is what keeps the key from naming a batch the graph was not built with -- the two directions + // bucket differently, and the graph builders read the same field for the same pass. if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_packed_ragged_graph) { cache_cfg.num_tokens_q = 0; cache_cfg.num_tokens_kv = 0; - cache_cfg.batch_size = static_cast(graph_dims(*this, pass).batch_size); + cache_cfg.batch_size = + pass == Pass::Fwd ? cache_cfg.graph_batch_size_fwd : cache_cfg.graph_batch_size_bwd; } // attn_scale is a pass-by-value graph input and different scales can share the same cached graph @@ -226,6 +240,47 @@ FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { return cache_cfg; } +std::string FusedAttnConfig::key_debug_string() const { + // Enums and sizes are printed as int64_t rather than by name, since the point is diffing two + // lines rather than reading one, and a numeric field cannot drift from a names table. + char buf[1024]; + std::snprintf( + buf, sizeof(buf), + "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, + static_cast(is_training), static_cast(deterministic), static_cast(cuda_graph), + static_cast(return_max_logit), static_cast(check_for_forward_support), + static_cast(attn_mask_type), static_cast(bias_type), + static_cast(window_size_left), static_cast(window_size_right), + static_cast(bottom_right_diagonal), static_cast(softmax_type), + static_cast(scaling_mode), static_cast(dropout), + static_cast(attn_scale), static_cast(qkv_dtype), + static_cast(o_dtype), static_cast(do_dtype), + static_cast(dqkv_dtype), static_cast(qkv_layout), + static_cast(o_format), static_cast(do_format), + static_cast(dqkv_layout), static_cast(qkv_scale_inv_format), + static_cast(do_scale_inv_format), static_cast(batch_size), + static_cast(num_attn_heads), static_cast(num_gqa_groups), + static_cast(head_dim_qk), static_cast(head_dim_v), + static_cast(max_seqlen_q), static_cast(max_seqlen_kv), + static_cast(num_tokens_q), static_cast(num_tokens_kv), + static_cast(bucketed_batch_size), static_cast(bucketed_num_tokens_q), + static_cast(bucketed_num_tokens_kv), static_cast(num_pages_k), + static_cast(num_pages_v), static_cast(page_size_k), + static_cast(page_size_v), static_cast(max_pages_per_seq_k), + static_cast(max_pages_per_seq_v), static_cast(bias_batch_size), + static_cast(bias_num_heads), static_cast(bias_seqlen_q), + static_cast(bias_seqlen_kv)); + return std::string(buf); +} + FusedAttnConfig FusedAttnFwdParams::make_config() const { const FusedAttnFwdParams ¶ms = *this; FusedAttnConfig cfg{}; diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 85f582bd19..5696e1e23b 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -11,6 +11,7 @@ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ +#include #include #include "common/common.h" @@ -20,10 +21,18 @@ namespace transformer_engine { namespace fused_attn { -// Which of the two graphs a config is being turned into. Declared here, rather than with the -// diagnostics that also name it, because two things about a config depend on the direction: -// make_cache_key() below, and graph_dims() further down. +// The pair that names one build site: whose graphs, and which of the two a config is being turned +// into. A site keeps its own graph cache and its own counters, so both halves travel together. // +// Declared here, with the config, rather than with the cache or its diagnostics: they are the +// vocabulary those two share, and the config is where their consumers start -- make_cache_key() +// below turns a config into one site's key, and derive() fills the dimensions built from it. +// +// Backend::F16 is the arbitrary-seqlen backend; the max512 one keeps no graph cache, so it has no +// site here. Narrower than the public NVTE_Fused_Attn_Backend, and not a substitute for it: this +// names only the backends that build graphs. +enum class Backend { F16, FP8 }; + // Passed in rather than derived, because a config cannot say which graph is being built from it. // check_for_forward_support and check_for_backward_support state which directions a caller wants // probed, and a backend query arriving from a framework has both set, so they answer a different @@ -106,15 +115,19 @@ struct FusedAttnConfig { // derive() recomputes unconditionally, so a config whose inputs change can simply be re-derived. bool is_derived = false; // THD batch/token counts, the raw buckets. The graph dimensions built out of them are - // graph_max_seqlen_* below and, because the batch is direction-dependent, graph_dims(). + // graph_max_seqlen_* below and, because the batch is direction-dependent, graph_batch_size_*. 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. + // Uses cu_seqlens or actual_seqlens. One answer per backend, because the same question has two: + // the FP8 graphs need newer cuDNN and frontend versions for it than the F16 ones, so a config + // that can hand cu_seqlens straight to one cannot necessarily hand them to the other. Each + // backend reads its own and no more; nothing reads both. bool uses_cu_seqlens_directly = false; + bool fp8_uses_cu_seqlens_directly = false; // Whether a ragged (THD) graph is built at packed token-count dimensions with ragged Stats/LSE, // rather than at dense max_seqlen ones. Held here rather than asked for at each of the places - // that need it -- graph_max_seqlen_* below, and graph_dims() -- because the key and the + // that need it -- graph_max_seqlen_* and graph_batch_size_* below -- because the key and the // graph have to be built at the same dimensions, and two independent queries are two chances to // disagree. Unlike the flags above, this one depends on the device as // well as the cuDNN version, so a config carries the answer for the device it was derived on; @@ -131,10 +144,22 @@ struct FusedAttnConfig { // dimensions the graph was built with -- a key that says otherwise is a hit on a graph of the // wrong shape -- and stating the substitution once is what keeps make_cache_key() and the graph // builders from drifting. Both passes build at the same sequence lengths; the batch size is the - // one dimension they disagree on, which is why that one comes from graph_dims() below instead of - // a field here. + // one dimension they disagree on, which is why that one is a pair below rather than a single + // field here. size_t graph_max_seqlen_q = 0; size_t graph_max_seqlen_kv = 0; + // The batch size the graph is built at and the width it expects ragged (THD) offsets in, one of + // each per direction. Pairs rather than one value apiece because the passes disagree on both, + // and a config is derived once and then probed and built for either direction, so no single + // field could answer: the selector derives a config and asks about forward and backward off that + // one copy. Whoever reads them names the direction they are building or keying for -- the graph + // builders, the code that binds runtime pointers to the built graph, and make_cache_key(), all + // of which have to agree, since a disagreement is a graph whose bound pointers do not describe + // the dimensions it was built at. All four are set together by the one condition in derive(). + size_t graph_batch_size_fwd = 0; + size_t graph_batch_size_bwd = 0; + DType ragged_offset_type_fwd = DType::kInt32; + DType ragged_offset_type_bwd = DType::kInt32; // Elements per token for each ragged tensor, from the layout group and the head dimensions. // Shared with the cu_seqlens_padded_to_offsets kernel, so the offsets the graph is told to // expect and the offsets that are written cannot drift apart. @@ -148,6 +173,28 @@ struct FusedAttnConfig { bool is_padding = false; bool is_causal = false; bool is_causal_bottom_right = false; + bool is_bias = false; + bool is_alibi = false; + bool is_softmax_offset = false; + bool is_mxfp8 = false; + // Whether the graph has a dropout node. The is_training term is what makes this one worth having + // as a field: a backward graph is only ever built for training, so the two directions used to + // spell this differently -- forward with the term, backward without -- and agreed only because + // every config that reaches a backward build has is_training set. One field states the rule the + // forward way, which is the safe way round: if that ever stops holding, a backward graph loses + // its dropout node rather than gaining one the cache key does not name. + bool is_dropout = false; + // Whether what each pass stores is itself quantized: O for a forward graph, dQKV for a backward + // one. Only the FP8 backend asks, and for it this is the whole of what separates the two + // tensor-scaling recipes -- FP8 out means the scale is known before the graph is built, F16 out + // means the graph has to compute it. Derived rather than asked at each build site so that the + // pairing of a pass with the tensor it writes is stated once. + // + // The FP8 builders read "not FP8" as "F16", which holds only because + // nvte_get_fused_attn_backend_v2 refuses an FP8 config whose output is neither before any graph + // is built; that refusal and these two fields are the same rule read from its two ends. + bool o_is_fp8 = false; + bool dqkv_is_fp8 = false; static constexpr size_t attr_sizes[] = { // basic attention settings @@ -260,6 +307,22 @@ struct FusedAttnConfig { // which fields are dropped, and whether the batch is bucketed -- so a key built for one pass // cannot be handed to the other's cache. FusedAttnConfig make_cache_key(Pass pass) const; + + // One line for the graph cache's level-2 trace: every field operator< compares, in its order, + // less device_id, which the trace's own prefix prints as the dev column. Abbreviated and terse + // on purpose, the reason to print a key at all being that two of these lines diff cleanly, + // naming the fields that cost an extra graph build. + // + // Four fields it prints that operator< does not compare, because a line without them cannot + // account for the values beside them: check_for_forward_support, which make_cache_key() sets + // from the pass and so is what says which direction's key this is, and the three bucketed_* + // inputs, which are what the normalization substituted into batch_size and the token counts. + // + // Defined here rather than with the diagnostics that print it because it is the third + // enumeration of these fields, after attr_sizes and operator< above. A field added to the key + // without being added here does not fail to build, it just stops appearing in the trace, so the + // three lists are kept where one change can see all of them. + std::string key_debug_string() const; }; // Assert that `cfg` has been through derive(), for code about to read a derived field. Worth @@ -277,32 +340,6 @@ inline void check_derived(const FusedAttnConfig &cfg) { "fused_attn.cpp."); } -// What a graph is built with that depends on which direction it is for, and so cannot be fields -// derive() fills: one stored value cannot answer for both passes, and the selector derives a config -// once and probes both directions off it. -struct GraphDims { - // The batch size the graph is built at: the bucketed one where a ragged layout is packed, so - // that one graph serves every batch in its bucket, and the true one otherwise. - int64_t batch_size = 0; - // The width the graph expects ragged (THD) offsets in. - DType ragged_offset_type = DType::kInt32; -}; - -// Both of the above for one direction. One function returning the pair rather than two returning -// one each, because a single condition decides both -- see the body -- and stating that condition -// twice is what would let the two drift into a graph built at a bucketed batch that reads its -// offsets at the other width. -// -// A free function rather than a member for the same reason as check_derived() above: it reads -// public fields and computes, and is no part of the config's own invariants. Requires a derived -// config and asserts it. -// -// Everything that has to agree about a graph asks this with the same pass: the builder, the code -// that binds runtime pointers to the built graph, and make_cache_key(). That is the point of -// having one place to ask, since a disagreement means a graph built at dimensions the pointers -// bound to it do not describe. -GraphDims graph_dims(const FusedAttnConfig &cfg, Pass pass); - 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/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index bd016e013e..a3aea68a5c 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -343,6 +343,38 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi const bool is_f16_or_bf16 = (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); + // Ask `verdict` about each direction the caller wants, and report the first refusal: the empty + // string means every direction asked about is served. Stated once here because every rule that + // is direction-dependent has to be asked this same way, and two copies of the gating would be + // two chances to probe a direction the caller never asked about. + // + // Forward is asked first because a config that cannot run forward cannot train either, and the + // forward refusal is the more useful of the two to report. Backward is skipped for inference, + // where no backward graph is ever built. + auto each_pass = [&](auto &&verdict) -> std::string { + if (cfg.check_for_forward_support) { + std::string reason = verdict(Pass::Fwd); + if (!reason.empty()) return reason; + } + if (cfg.is_training && cfg.check_for_backward_support) { + std::string reason = verdict(Pass::Bwd); + if (!reason.empty()) return reason; + } + return ""; + }; + + // cuDNN's own verdict on `backend`. The two backends differ only in which set of graphs gets + // built, so they share this; what is theirs alone are the rules in each branch below. + // + // Each backend answers for both directions from its own translation unit, the only place that can + // name the graph builders, so choosing between them is all the dispatch left to do here. + auto probe = [&](Backend backend) -> std::string { + return each_pass([&](Pass pass) { + return backend == Backend::FP8 ? support_verdict_fp8(cfg, pass, handle) + : support_verdict_f16(cfg, pass, handle); + }); + }; + if (is_fp8) { if (cfg.return_max_logit) { return reject(message, "FP8 fused attention does not support return_max_logit=True."); @@ -352,14 +384,50 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi return reject(message, "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + std::to_string(static_cast(qkv_format)) + "."); } - if (cfg.check_for_forward_support) { - std::string fwd_reason = is_supported_fp8_fwd(cfg, handle); - if (!fwd_reason.empty()) return reject(message, std::move(fwd_reason)); + // The rest of what the FP8 graphs cannot represent: bias, ALiBi, and the quantization recipes + // they are not written for. TE's rules rather than cuDNN's, and stated here rather than in the + // build path for the reason all the rules above are: a rejection stated here is an answer + // carrying its reason, where the same rule inside a graph build would have to travel out as an + // exception. + if (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { + return reject(message, "FP8 fused attention does not support pre/post_scale_bias yet!"); } - if (cfg.is_training && cfg.check_for_backward_support) { - std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); - if (!bwd_reason.empty()) return reject(message, std::move(bwd_reason)); + if (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI) { + return reject(message, "FP8 fused attention does not support ALiBi yet!"); + } + + // Whether the config names a recipe the FP8 graphs are written for at all. Delayed scaling + // writes FP8 out and keeps its scale, current scaling writes F16 and computes one, MXFP8 writes + // F16 with block scales; every other pairing of scaling mode and output dtype is refused here. + // + // Per direction, because the pairing is read off what each pass stores, which is also how the + // graph builders read which of the three they are building for -- off cfg.o_is_fp8 and + // cfg.dqkv_is_fp8, taking "not FP8" to mean F16. That reading is sound only because this + // refusal has already happened, so the two belong to each other: a pairing accepted here must + // be one they read the same way, and anything added to either belongs in both. + std::string recipe_reason = each_pass([&](Pass pass) -> std::string { + const NVTEDType out_dtype = (pass == Pass::Fwd) ? cfg.o_dtype : cfg.dqkv_dtype; + const bool out_is_fp8 = (pass == Pass::Fwd) ? cfg.o_is_fp8 : cfg.dqkv_is_fp8; + const bool out_is_f16 = (out_dtype == kNVTEFloat16 || out_dtype == kNVTEBFloat16); + const bool serves_this_output = + (cfg.scaling_mode == NVTE_DELAYED_TENSOR_SCALING && (out_is_fp8 || out_is_f16)) || + (cfg.scaling_mode == NVTE_MXFP8_1D_SCALING && out_is_f16); + if (!serves_this_output) { + return "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 " + "recipes!"; + } + return ""; + }); + if (!recipe_reason.empty()) return reject(message, std::move(recipe_reason)); + + // Asked after the pairing above, not with it: a config that names no recipe at all should hear + // that rather than be sent to upgrade cuDNN for a recipe it was not asking for. + if (cfg.scaling_mode == NVTE_MXFP8_1D_SCALING && cudnn_runtime_version < 92100) { + return reject(message, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); } + + std::string reason = probe(Backend::FP8); + if (!reason.empty()) return reject(message, std::move(reason)); return NVTE_Fused_Attn_Backend::NVTE_FP8; } @@ -378,14 +446,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { return reject(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); } - if (cfg.check_for_forward_support) { - std::string fwd_reason = is_supported_f16_fwd(cfg, handle); - if (!fwd_reason.empty()) return reject(message, std::move(fwd_reason)); - } - if (cfg.is_training && cfg.check_for_backward_support) { - std::string bwd_reason = is_supported_f16_bwd(cfg, handle); - if (!bwd_reason.empty()) return reject(message, std::move(bwd_reason)); - } + std::string reason = probe(Backend::F16); + if (!reason.empty()) return reject(message, std::move(reason)); return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } @@ -445,9 +507,10 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // support query and the execution path reach the same cache through the same accessor. That is // what the HIT below means: by the time a backend has been selected, the entry the implementation // needs has already been built and inserted by the probe that selected it, so what was checked is -// what runs. The rules the selector does state for itself are the ones cuDNN cannot answer, -// because they are about whether the graph computes what was asked for rather than whether cuDNN -// can run it. +// what runs. The rules the selector does state for itself are the ones cuDNN cannot answer: either +// about whether the graph computes what was asked for rather than whether cuDNN can run it, or +// about what TE's graphs can represent in the first place, which is where the FP8 recipes come in. +// Stating them here rather than inside a build is what lets each one answer with its reason. // // nvte_fused_attn_fwd_v2 // | @@ -456,17 +519,19 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // +-- nvte_get_fused_attn_backend_v2 the support query // | | // | +-- TE's own rules: THD and paged KV need a padding mask, no pre-scale bias, -// | | ragged Q/KV needs sm90+, the cuDNN 9.15-and-older CUDA-graph bug +// | | ragged Q/KV needs sm90+, the cuDNN 9.15-and-older CUDA-graph bug, and for FP8 +// | | bias, ALiBi and the quantization recipes its graphs are not written for // | | `-- reject -> NVTE_No_Backend + reason -> the NVTE_ERROR below // | | -// | `-- is_supported_f16_fwd / is_supported_fp8_fwd -// | `-- cache_graph_f16_fwd(): builds and inserts the entry, or throws, in -// | which case cuDNN's message becomes the reason for the refusal +// | `-- probe(backend) -> support_verdict_f16 / support_verdict_fp8, with Pass::Fwd +// | `-- support_verdict() +// | `-- get_graph(): builds and inserts the entry, or throws, +// | in which case cuDNN's message becomes the reason for the refusal // | // `-- fused_attn_arbitrary_seqlen_fwd -> ..._fwd_impl the selected backend // | -// +-- cache_graph_f16_fwd() HIT: the entry the query above just built -// +-- build_plans() the kernel compilation, once per entry +// +-- get_graph() HIT: the entry the query above just built +// +-- build_plans() the kernel compilation, once per entry // `-- bind device pointers, graph.execute() void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { NVTE_API_CALL(nvte_fused_attn_fwd_v2); 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 63094fa8c4..e68595fd26 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 @@ -25,12 +25,7 @@ namespace fused_attn { namespace fe = cudnn_frontend; -// Every graph-cache event raised here names the build site it came from. This file is the f16 -// arbitrary-seqlen backend throughout; only the pass differs between call sites. Pass itself needs -// no using-declaration: it is fused_attn::Pass, since the config answers by direction too. -using graph_cache_debug::Backend; - -using SdpaF16FwdGraphAndTensors = +using F16FwdGraphAndTensors = std::tuple, std::shared_ptr, // Q std::shared_ptr, // K @@ -54,18 +49,17 @@ using SdpaF16FwdGraphAndTensors = std::shared_ptr>; // dropout_offset // Constructs the forward graph for one cache key, and only constructs it: whether cuDNN will run -// it is settled by the caller, in lookup_or_cache_graph(), which is also where the plan build -// eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every call -// that does need one now sits on the other side of that boundary. +// it is settled by the caller, in cache_graph(), which is also where the plan build eventually +// happens. Hence no cuDNN handle here -- describing a graph needs none, and every call that does +// need one now sits on the other side of that boundary. // // Everything the graph's shape and topology depends on comes from `cfg`, so the build has one // source of truth and cannot drift from the caller that will bind pointers to it. The two -// dimensions the config cannot answer on its own -- the batch size, and the width ragged offsets -// are written in, both of which differ between the passes -- come from graph_dims() asked with -// Pass::Fwd, the same way the code binding pointers to this graph asks. -static SdpaF16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg) { - const GraphDims dims = graph_dims(cfg, Pass::Fwd); - const int64_t b = dims.batch_size; +// dimensions that differ between the passes -- the batch size, and the width ragged offsets are +// written in -- are read from the config's forward halves, the same ones the code binding pointers +// to this graph reads. +static F16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg) { + const int64_t b = static_cast(cfg.graph_batch_size_fwd); const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); const cudnn_frontend::DataType_t tensorType = @@ -86,29 +80,24 @@ static SdpaF16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); const int64_t window_size_left = cfg.window_size_left; const int64_t window_size_right = cfg.window_size_right; - const bool is_training = cfg.is_training; const bool return_max_logit = cfg.return_max_logit; const float dropout_probability = cfg.dropout; const NVTE_QKV_Layout qkv_layout = cfg.qkv_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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; - const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - const bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_bias = cfg.is_bias; + const bool is_alibi = cfg.is_alibi; + const bool is_causal = cfg.is_causal; const bool is_causal_bottom_right = cfg.is_causal_bottom_right; const bool is_padding = cfg.is_padding; const bool is_paged_kv = cfg.is_paged_kv; - const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const bool is_dropout = (is_training && dropout_probability != 0.0f); + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; const auto cudnn_runtime_version = cudnnGetVersion(); const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = dims.ragged_offset_type; + const DType ragged_offset_type = cfg.ragged_offset_type_fwd; const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; const bool generate_stats = true; // Always return stats @@ -373,19 +362,6 @@ static SdpaF16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg offset_kv_tuple, offset_s_tuple, dropout_tuple); } -// The forward graph cache and the only route to it. Both the execution path and the support -// probe come through here, so a probe leaves behind exactly the entry a later execution finds. -// That is what lets the probe's answer describe the graph that actually runs, rather than a -// separately built lookalike. -static std::shared_ptr> cache_graph_f16_fwd( - const FusedAttnConfig &cfg, cudnnHandle_t handle) { - static GraphCache cache; - // Asserted once here for both the key and the graph, which read the same derived fields. - check_derived(cfg); - return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Fwd), Backend::F16, Pass::Fwd, - handle, [&] { return create_graph_f16_fwd(cfg); }); -} - void fused_attn_arbitrary_seqlen_fwd_impl( const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrK, void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, void *devPtrO, @@ -395,11 +371,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - // Asked with the same pass the graph was built with, so that the dimensions below and the ones - // the graph was built at cannot be decided differently. - const GraphDims dims = graph_dims(cfg, Pass::Fwd); - const int64_t b = dims.batch_size; - const DType ragged_offset_type = dims.ragged_offset_type; + // Read from the same halves of the config the graph was built from, so that the dimensions below + // and the ones the graph was built at cannot be decided differently. Asserted derived here + // because these are the first derived fields this path reads, ahead of the get_graph() that + // asserts it for the build. + check_derived(cfg); + const int64_t b = static_cast(cfg.graph_batch_size_fwd); + const DType ragged_offset_type = cfg.ragged_offset_type_fwd; // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever the // bucketing above did to `b`. const int64_t actual_b = static_cast(cfg.batch_size); @@ -409,10 +387,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const bool return_max_logit = cfg.return_max_logit; // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; - const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; const bool is_paged_kv = cfg.is_paged_kv; @@ -422,10 +400,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; try { - auto cache_entry = cache_graph_f16_fwd(cfg, handle); + auto cache_entry = get_graph(cfg, handle); 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] = cache_entry->tensors; + dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; // This graph is going to be used, so finish the build the cache deferred. build_plans(Backend::F16, Pass::Fwd, *cache_entry); @@ -457,8 +435,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - graph_cache_debug::record_exec(Backend::F16, Pass::Fwd); - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -567,12 +543,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::F16, Pass::Fwd); } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } } -using SdpaF16BwdGraphAndTensors = +using F16BwdGraphAndTensors = std::tuple, std::shared_ptr, // q std::shared_ptr, // k @@ -599,10 +576,9 @@ using SdpaF16BwdGraphAndTensors = std::shared_ptr>; // dropout_offset // The backward counterpart of create_graph_f16_fwd; see there for why it constructs the graph and -// nothing else, and why the two direction-dependent dimensions are asked for rather than stored. -static SdpaF16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg) { - const GraphDims dims = graph_dims(cfg, Pass::Bwd); - const int64_t b = dims.batch_size; +// nothing else, and why the two direction-dependent dimensions come from the config in pairs. +static F16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg) { + const int64_t b = static_cast(cfg.graph_batch_size_bwd); const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); const cudnn_frontend::DataType_t tensorType = @@ -619,25 +595,21 @@ static SdpaF16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg const int64_t window_size_right = cfg.window_size_right; const float dropout_probability = cfg.dropout; const NVTE_QKV_Layout qkv_layout = cfg.qkv_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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; const bool deterministic = cfg.deterministic; - const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - const bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_bias = cfg.is_bias; + const bool is_alibi = cfg.is_alibi; + const bool is_causal = cfg.is_causal; const bool is_causal_bottom_right = cfg.is_causal_bottom_right; const bool is_padding = cfg.is_padding; - const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const bool is_dropout = (dropout_probability != 0.0f); + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); const bool use_packed_ragged_graph = cfg.uses_packed_ragged_graph; const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = dims.ragged_offset_type; + const DType ragged_offset_type = cfg.ragged_offset_type_bwd; auto mha_graph = std::make_shared(); mha_graph->set_io_data_type(tensorType) @@ -867,15 +839,6 @@ static SdpaF16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg offset_s_tuple, dropout_tuple); } -// The backward counterpart of cache_graph_f16_fwd; see there. -static std::shared_ptr> cache_graph_f16_bwd( - const FusedAttnConfig &cfg, cudnnHandle_t handle) { - static GraphCache cache; - check_derived(cfg); - return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Bwd), Backend::F16, Pass::Bwd, - handle, [&] { return create_graph_f16_bwd(cfg); }); -} - void fused_attn_arbitrary_seqlen_bwd_impl( const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrKTranspose, void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, void *devPtrSoftmaxOffset, @@ -886,29 +849,31 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cudnnHandle_t handle) { using namespace transformer_engine; - // Asked with the same pass the graph was built with, so that the dimensions below and the ones - // the graph was built at cannot be decided differently. - const GraphDims dims = graph_dims(cfg, Pass::Bwd); - const int64_t b = dims.batch_size; - const DType ragged_offset_type = dims.ragged_offset_type; + // Read from the same halves of the config the graph was built from, so that the dimensions below + // and the ones the graph was built at cannot be decided differently. Asserted derived here + // because these are the first derived fields this path reads, ahead of the get_graph() that + // asserts it for the build. + check_derived(cfg); + const int64_t b = static_cast(cfg.graph_batch_size_bwd); + const DType ragged_offset_type = cfg.ragged_offset_type_bwd; // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by. const int64_t actual_b = static_cast(cfg.batch_size); const bool use_ragged_stats = cfg.uses_ragged_stats; // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; - const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const bool is_dropout = (cfg.dropout != 0.0f); + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; const bool is_ragged_q = cfg.is_ragged_q; const bool is_ragged_kv = cfg.is_ragged_kv; try { - auto cache_entry = cache_graph_f16_bwd(cfg, handle); + auto cache_entry = get_graph(cfg, handle); 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] = cache_entry->tensors; + dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; // This graph is going to be used, so finish the build the cache deferred. build_plans(Backend::F16, Pass::Bwd, *cache_entry); @@ -935,8 +900,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - graph_cache_debug::record_exec(Backend::F16, Pass::Bwd); - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -1029,6 +992,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::F16, Pass::Bwd); } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } @@ -1245,30 +1209,17 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i } } -// Whether cuDNN can run the forward graph this config asks for: the empty string if it can, -// otherwise cuDNN's own account of why not, which the backend selector reports to the caller. +// The one entry point into this translation unit's support probes; see fused_attn::support_verdict, +// which is all it does. It exists because create_graph_f16_* is file-local, so this is the only +// place that can name it, and because the selector calls it from another translation unit. // -// The question is answered by deriving the graph's inputs and building the graph, which is -// where every rejection comes from -- there is no separate list of rules to keep in step with -// the builder. The graph goes into the same cache the execution path reads, so the work is not -// thrown away and what was checked is what will run. It stops short of graph.build_plans(), the -// expensive step, which the first execution of the graph does instead; see CachedGraph. -// -// A refusal, by contrast, is not cached: nothing is stored for a key cuDNN rejected, so asking the -// same question again pays for the build again. See lookup_or_cache_graph. -// -// The direction comes from which of these two functions was called, not from the config: a config -// arriving from a framework has both check_for_*_support set, so both probes run off a single -// config, and each has to name its own direction for the key and the graph to be the forward ones. -std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - return fused_attn::support_verdict("is_supported_f16_fwd", - [&] { fused_attn::cache_graph_f16_fwd(cfg, handle); }); -} - -// The backward counterpart of is_supported_f16_fwd; see there. -std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - return fused_attn::support_verdict("is_supported_f16_bwd", - [&] { fused_attn::cache_graph_f16_bwd(cfg, handle); }); +// Turning `pass` into the template argument is the whole of the body: the direction has to be a +// compile-time constant to pick a builder, and this is where the two meet. +std::string support_verdict_f16(const FusedAttnConfig &cfg, Pass pass, cudnnHandle_t handle) { + if (pass == Pass::Fwd) { + return fused_attn::support_verdict(cfg, handle); + } + return fused_attn::support_verdict(cfg, handle); } } // 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 c493b5ee1b..a373bb5e85 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 @@ -41,15 +41,14 @@ void fused_attn_arbitrary_seqlen_bwd(const fused_attn::FusedAttnConfig &cfg, con const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); -// check if a given configuration is supported for F16/BF16 forward; -// if it is, cache the graph built for this config, and return an empty string; -// if not, return a diagnostic message explaining why it is not supported. -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 fused_attn::FusedAttnConfig &cfg, cudnnHandle_t handle); +// cuDNN's verdict on this config's F16/BF16 graph for `pass`: an empty string if it can run, +// otherwise a diagnostic message explaining why not. A verdict of "supported" leaves the graph in +// the cache, where the execution path finds it. +// +// The direction is a runtime argument, not two functions, because the graph builder it selects is +// local to this translation unit -- so this is the only place that can map one to the other. +std::string support_verdict_f16(const fused_attn::FusedAttnConfig &cfg, fused_attn::Pass pass, + 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 8df1e3bfb2..1c402f2676 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -20,13 +20,8 @@ namespace fused_attn { using namespace transformer_engine; namespace fe = cudnn_frontend; -// Every graph-cache event raised here names the build site it came from. This file is the fp8 -// backend throughout; only the pass differs between call sites. Pass itself needs no -// using-declaration: it is fused_attn::Pass, since the config answers by direction too. -using graph_cache_debug::Backend; - // fused attention FWD FP8 with FE 1.0+ -using SdpaFp8FwdGraphAndTensors = +using Fp8FwdGraphAndTensors = std::tuple, std::shared_ptr, // Q std::shared_ptr, // K @@ -49,99 +44,32 @@ using SdpaFp8FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The FP8 policy decisions a graph is built from, which the config cannot state on its own: which -// quantization recipe the graph quantizes for, whether cu_seqlens can be handed to cuDNN directly, -// and whether O arrives in F16. Each decides which tensors the graph has, and so which pointers -// the variant pack has to bind, so the build and the execution ask the same question here rather -// than each deciding for itself. Free functions for the reason check_derived() is one: they read -// public fields and compute, and belong to this backend rather than to the config. +// The three recipes these graphs are written for, spelled the same way at each of the four sites +// that build or bind one: +// +// is_mxfp8 = scaling_mode is MXFP8 +// is_delayed_scaling = !is_mxfp8 && +// is_current_scaling = !is_mxfp8 && ! +// +// which is exactly one of the three by construction, no combination of the booleans being able to +// say two things at once. The output half comes from cfg.o_is_fp8 or cfg.dqkv_is_fp8 -- a forward +// graph writes O, a backward one dQKV -- and reads "not FP8" as F16, which holds because +// nvte_get_fused_attn_backend_v2 refuses an FP8 config whose output is neither. See there for the +// rest of what these graphs cannot represent, and config_and_params.h for the two fields. // // Unlike the F16 path there is no bucketing to do, because FP8 has no ragged/THD support: the // graph's shapes are exactly the config's. -// Which quantization recipe the graph is built for. One enum rather than the three mutually -// exclusive booleans it replaces, since exactly one recipe applies to a config and a triple leaves -// the other seven combinations expressible. -enum class Fp8Recipe { DelayedScaling, CurrentScaling, MxFp8 }; - -// The recipe `cfg` asks for, or a throw naming what FP8 cannot serve. The rejections are TE's own -// rather than cuDNN's -- bias, ALiBi and the recipe combinations -- and a support probe reports -// them the same way it reports a cuDNN refusal, as the reason the FP8 backend was not selected. -// -// The pass decides which tensor the recipe is read off: the forward graph writes O and the -// backward writes dQKV, and a run can quantize one without the other, so each pass reads the -// dtype of what it actually stores. -static Fp8Recipe fp8_recipe(const FusedAttnConfig& cfg, Pass pass) { - check_derived(cfg); - const cudnn_frontend::DataType_t out_type = - get_cudnn_fe_dtype(static_cast(pass == Pass::Fwd ? cfg.o_dtype : cfg.dqkv_dtype)); - const NVTEScalingMode scaling_mode = cfg.scaling_mode; - const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - const bool is_alibi = (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI); - - 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!"); - const bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (out_type == cudnn_frontend::DataType_t::FP8_E4M3 || - out_type == cudnn_frontend::DataType_t::FP8_E5M2); - const bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (out_type == cudnn_frontend::DataType_t::HALF || - out_type == cudnn_frontend::DataType_t::BFLOAT16); - const bool is_mxfp8 = - (scaling_mode == NVTE_MXFP8_1D_SCALING) && (out_type == cudnn_frontend::DataType_t::HALF || - out_type == cudnn_frontend::DataType_t::BFLOAT16); - NVTE_CHECK( - is_delayed_scaling || is_current_scaling || is_mxfp8, - "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); - NVTE_CHECK(!is_mxfp8 || cudnnGetVersion() >= 92100, - "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - - if (is_delayed_scaling) return Fp8Recipe::DelayedScaling; - if (is_current_scaling) return Fp8Recipe::CurrentScaling; - return Fp8Recipe::MxFp8; -} - -// Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative -// tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike -// the F16 path, the FP8 path has no THD/ragged-offset support, so only the -// cu_seqlens_to_actual_seqlens conversion applies here. Also note that the -// needed versions of cuDNN backend and frontend are higher than for F16, which is why this does -// not read cfg.uses_cu_seqlens_directly, the F16 path's answer to the same question.) -static bool fp8_uses_cu_seqlens_directly(const FusedAttnConfig& cfg) { - const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); - return - // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). - // Note: For the Python API, 1.27 is required. - CUDNN_FRONTEND_VERSION >= 12600 && - // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN - // version, so we'll do the same. - (CUDNN_VERSION >= 92500 && cudnnGetVersion() >= 92500) && - // 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; -} - -// Whether O arrives in F16 rather than FP8, which is what decides if the backward graph has to -// descale it on the way in. Read off O for both passes, unlike the recipe. -static bool fp8_o_in_f16(const FusedAttnConfig& cfg) { - const cudnn_frontend::DataType_t o_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); - return o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16; -} - // Constructs the forward FP8 graph for one cache key, and only constructs it: whether cuDNN will -// run it is settled by the caller, in lookup_or_cache_graph(), which is also where the plan -// build eventually happens. Hence no cuDNN handle here -- describing a graph needs none, and every -// call that does need one now sits on the other side of that boundary. +// run it is settled by the caller, in cache_graph(), which is also where the plan build eventually +// happens. Hence no cuDNN handle here -- describing a graph needs none, and every call that does +// need one now sits on the other side of that boundary. // // Everything the graph's shape and topology depends on comes from `cfg`, so the build has one -// source of truth and cannot drift from the caller that will bind pointers to it. The decisions -// the config cannot state itself are asked for with Pass::Fwd, the same way the code binding -// pointers to this graph asks. -static SdpaFp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { +// source of truth and cannot drift from the caller that will bind pointers to it -- including the +// forward half of the pass-indexed fields, read here the same way the code binding pointers to +// this graph reads it. +static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t qkv_tensor_type = get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); @@ -161,22 +89,17 @@ static SdpaFp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg 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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; - const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_bias = cfg.is_bias; + const bool is_causal = cfg.is_causal; const bool is_causal_bottom_right = cfg.is_causal_bottom_right; const bool is_padding = cfg.is_padding; - const bool is_dropout = (is_training && dropout_probability != 0.0f); - const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Fwd); - const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; - const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; - const bool is_mxfp8 = recipe == Fp8Recipe::MxFp8; - const bool use_cu_seqlens_directly = fp8_uses_cu_seqlens_directly(cfg); + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_delayed_scaling = !is_mxfp8 && cfg.o_is_fp8; + const bool is_current_scaling = !is_mxfp8 && !cfg.o_is_fp8; + const bool use_cu_seqlens_directly = cfg.fp8_uses_cu_seqlens_directly; auto mha_graph = std::make_shared(); mha_graph->set_io_data_type(qkv_tensor_type) @@ -420,17 +343,6 @@ static SdpaFp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg softmax_offset_tuple, padding_tuple, dropout_tuple); } -// The FP8 forward graph cache and the only route to it. Both the execution path and the support -// probe come through here, so a probe leaves behind exactly the entry a later execution finds. -static std::shared_ptr> cache_graph_fp8_fwd( - const FusedAttnConfig& cfg, cudnnHandle_t handle) { - static GraphCache cache; - // Asserted once here for both the key and the graph, which read the same derived fields. - check_derived(cfg); - return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Fwd), Backend::FP8, Pass::Fwd, - handle, [&] { return create_graph_fp8_fwd(cfg); }); -} - 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, @@ -441,27 +353,30 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - // Asked with the same pass the graph was built with, so that the tensors bound below and the - // ones the graph was built with cannot be decided differently. Also where an unserviceable - // configuration is rejected, ahead of the cache lookup. - const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Fwd); - const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; - const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; - const bool use_cu_seqlens_directly = fp8_uses_cu_seqlens_directly(cfg); + // Asserted derived here because the reads below are the first derived fields this path touches, + // ahead of the get_graph() that asserts it for the build. + check_derived(cfg); + + // Read from the same fields the graph was built from, so that the tensors bound below and the + // ones the graph was built with cannot be decided differently. + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_delayed_scaling = !is_mxfp8 && cfg.o_is_fp8; + const bool is_current_scaling = !is_mxfp8 && !cfg.o_is_fp8; + const bool use_cu_seqlens_directly = cfg.fp8_uses_cu_seqlens_directly; const int64_t b = static_cast(cfg.batch_size); // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; - const bool is_dropout = (cfg.is_training && cfg.dropout != 0.0f); - const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; try { - auto cache_entry = cache_graph_fp8_fwd(cfg, handle); + auto cache_entry = get_graph(cfg, handle); 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] = cache_entry->tensors; + dropout_offset] = cache_entry->graph_and_tensors; // This graph is going to be used, so finish the build the cache deferred. build_plans(Backend::FP8, Pass::Fwd, *cache_entry); @@ -476,8 +391,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::record_exec(Backend::FP8, Pass::Fwd); - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -538,13 +451,14 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::FP8, Pass::Fwd); } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } } // fused attention BWD FP8 with FE 1.0+ -using SdpaFp8BwdGraphAndTensors = +using Fp8BwdGraphAndTensors = std::tuple, std::shared_ptr, // Q std::shared_ptr, // Q_t @@ -589,14 +503,14 @@ using SdpaFp8BwdGraphAndTensors = std::shared_ptr>; // dropout_offset // Builds the backward FP8 graph for one cache key, up to check_support() but not -// graph.build_plans(); see CachedGraph for why the plan build is left to whoever runs the graph. +// graph.build_plans(); see CacheEntry for why the plan build is left to whoever runs the graph. // // Everything the graph's shape and topology depends on is re-derived from `cfg` here, so the // build has one source of truth for them. Unlike the F16 path, FP8 has no ragged/THD support, // so the shapes are exactly the config's and need no bucketing from the caller. // The backward counterpart of create_graph_fp8_fwd; see there for why it constructs the graph // and nothing else. -static SdpaFp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { +static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t qkv_tensor_type = get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); @@ -622,23 +536,20 @@ static SdpaFp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg const NVTE_QKV_Format do_format = cfg.do_format; 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 bool bottom_right_diagonal = cfg.bottom_right_diagonal; const bool deterministic = cfg.deterministic; - const bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - const bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + const bool is_bias = cfg.is_bias; + const bool is_causal = cfg.is_causal; const bool is_causal_bottom_right = cfg.is_causal_bottom_right; const bool is_padding = cfg.is_padding; - const bool is_dropout = (dropout_probability != 0.0f); - const bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Bwd); - const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; - const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; - const bool is_mxfp8 = recipe == Fp8Recipe::MxFp8; - const bool is_O_in_F16 = fp8_o_in_f16(cfg); + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_delayed_scaling = !is_mxfp8 && cfg.dqkv_is_fp8; + const bool is_current_scaling = !is_mxfp8 && !cfg.dqkv_is_fp8; + // Whether O arrived in F16 rather than FP8, which decides whether this graph has to descale it on + // the way in. Read off O, unlike the recipe above, because O is what the forward pass stored. + const bool is_O_in_F16 = !cfg.o_is_fp8; auto mha_graph = std::make_shared(); @@ -1013,15 +924,6 @@ static SdpaFp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); } -// The backward counterpart of cache_graph_fp8_fwd; see there. -static std::shared_ptr> cache_graph_fp8_bwd( - const FusedAttnConfig& cfg, cudnnHandle_t handle) { - static GraphCache cache; - check_derived(cfg); - return lookup_or_cache_graph(cache, cfg.make_cache_key(Pass::Bwd), Backend::FP8, Pass::Bwd, - handle, [&] { return create_graph_fp8_bwd(cfg); }); -} - void fused_attn_fp8_bwd_impl( const FusedAttnConfig& cfg, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, @@ -1036,31 +938,33 @@ void fused_attn_fp8_bwd_impl( cudnnHandle_t handle) { using namespace transformer_engine; - // Asked with the same pass the graph was built with, so that the tensors bound below and the - // ones the graph was built with cannot be decided differently. Also where an unserviceable - // configuration is rejected, ahead of the cache lookup. - const Fp8Recipe recipe = fp8_recipe(cfg, Pass::Bwd); - const bool is_delayed_scaling = recipe == Fp8Recipe::DelayedScaling; - const bool is_current_scaling = recipe == Fp8Recipe::CurrentScaling; - const bool is_mxfp8 = recipe == Fp8Recipe::MxFp8; - const bool is_O_in_F16 = fp8_o_in_f16(cfg); + // Asserted derived here because the reads below are the first derived fields this path touches, + // ahead of the get_graph() that asserts it for the build. + check_derived(cfg); + + // Read from the same fields the graph was built from, so that the tensors bound below and the + // ones the graph was built with cannot be decided differently. + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_delayed_scaling = !is_mxfp8 && cfg.dqkv_is_fp8; + const bool is_current_scaling = !is_mxfp8 && !cfg.dqkv_is_fp8; + const bool is_O_in_F16 = !cfg.o_is_fp8; const int64_t b = static_cast(cfg.batch_size); const int64_t h = static_cast(cfg.num_attn_heads); // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const bool is_bias = (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; - const bool is_dropout = (cfg.dropout != 0.0f); - const bool is_softmax_offset = (cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; try { - auto cache_entry = cache_graph_fp8_bwd(cfg, handle); + auto cache_entry = get_graph(cfg, handle); 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, 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] = cache_entry->tensors; + dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; // This graph is going to be used, so finish the build the cache deferred. build_plans(Backend::FP8, Pass::Bwd, *cache_entry); @@ -1073,8 +977,6 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } - graph_cache_debug::record_exec(Backend::FP8, Pass::Bwd); - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -1159,6 +1061,7 @@ void fused_attn_fp8_bwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::FP8, Pass::Bwd); } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } @@ -1392,29 +1295,16 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const } } -// Whether the FP8 forward graph this config asks for can run: the empty string if it can, -// otherwise the account of why not, which the backend selector reports to the caller. -// -// The question is answered by building the graph, which is where every rejection comes from -- -// there is no separate list of rules to keep in step with the builder. The graph goes into the same -// cache the execution path reads, so the work is not thrown away and what was checked is what will -// run. It stops short of graph.build_plans(), the expensive step, which the first execution of the -// graph does instead; see CachedGraph. +// The FP8 counterpart of support_verdict_f16; see there for why the direction arrives at runtime. // -// Unlike the F16 path, some of the rejections here are TE's own rather than cuDNN's: fp8_recipe() -// throws for bias, ALiBi and the recipe combinations FP8 does not serve, from inside the build. -// They read the same to the selector, which wants a reason and does not care whose rule it was. -// -// The direction comes from which of these two functions was called; see is_supported_f16_fwd. -std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { - return fused_attn::support_verdict("is_supported_fp8_fwd", - [&] { fused_attn::cache_graph_fp8_fwd(cfg, handle); }); -} - -// The backward counterpart of is_supported_fp8_fwd; see there. -std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { - return fused_attn::support_verdict("is_supported_fp8_bwd", - [&] { fused_attn::cache_graph_fp8_bwd(cfg, handle); }); +// Only cuDNN's rules reach this. TE's own -- bias, ALiBi and the recipes these graphs are not +// written for -- are stated in nvte_get_fused_attn_backend_v2 and answered before it ever gets +// here, which is why no build on this path throws for a configuration it cannot serve. +std::string support_verdict_fp8(const FusedAttnConfig& cfg, Pass pass, cudnnHandle_t handle) { + if (pass == Pass::Fwd) { + return fused_attn::support_verdict(cfg, handle); + } + return fused_attn::support_verdict(cfg, handle); } } // 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 9e8f997d98..2749aa1fc1 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -39,15 +39,9 @@ void fused_attn_fp8_bwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *in 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. -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 fused_attn::FusedAttnConfig &cfg, cudnnHandle_t handle); +// The FP8 counterpart of support_verdict_f16; see there. +std::string support_verdict_fp8(const fused_attn::FusedAttnConfig &cfg, fused_attn::Pass pass, + 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.h b/transformer_engine/common/fused_attn/graph_cache.h index d96d7166ca..9e7b344342 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -10,8 +10,9 @@ // // The four build sites (f16 and fp8, forward and backward) differ only in how // they construct their graph and which tensors they hand back. Everything after -// that -- the lookup, the locking, the support check, the once-per-entry plan -// build -- is shared, and lives here rather than in four copies. +// that -- the cache each one keeps, the lookup, the locking, the support check, +// the once-per-entry plan build -- is shared, and lives here rather than in four +// copies: a site names its backend, pass and graph builder to get_graph(). // // The five frontend calls a graph goes through, and which of our functions pays for // each. The frontend's are written graph.*, since that is how they are invoked and @@ -49,32 +50,6 @@ namespace transformer_engine { namespace fused_attn { -// The verdict an is_supported_* helper reports for `probe`: the empty string if it completes, -// otherwise cuDNN's own account of the refusal. Support is discovered by building the graph, so a -// probe is one call and everything else is what to do with a failure; this is that, once, for all -// four helpers. -// -// Refusals and failures on the way to a verdict read alike, because CUDNN_BACKEND_API_FAILED -- -// raised for any non-success cudnnStatus_t -- cannot separate CUDNN_STATUS_NOT_SUPPORTED from -// CUDNN_STATUS_ALLOC_FAILED. Either way this backend cannot serve this call, and either way what -// the caller wants is the message. -// -// `what` names the probe and is used only when a failure carried no message of its own: support is -// signalled by returning the empty string, so an empty refusal would read as an endorsement. -template -std::string support_verdict(const char *what, ProbeFn &&probe) { - try { - probe(); - return ""; - } catch (const std::exception &e) { - const char *reason = e.what(); - if (reason != nullptr && reason[0] != '\0') return reason; - return std::string(what) + ": rejected without a reason."; - } catch (...) { - return std::string(what) + ": unknown failure."; - } -} - // A graph in the cache, plus the tensor attributes needed to bind runtime pointers to it. // // Entries are built only as far as check_support(), which is all it takes to decide whether a @@ -82,17 +57,18 @@ std::string support_verdict(const char *what, ProbeFn &&probe) { // of the five frontend calls -- is left to the execution path, since a support query never runs the // graph and many of the keys it builds are never run by anything. // -// plans_built guards that completion, which has to happen exactly once per entry: the entry is +// build_plans_once guards that completion, which has to happen exactly once per entry: the entry is // shared across threads and graph.build_plans() mutates it in place. Keeping the flag in the entry // keeps it with the graph it describes and leaves unrelated keys free to build concurrently. A // build that throws leaves it unset, so a later call retries rather than executing a graph with no // plans. template -struct CachedGraph { - explicit CachedGraph(GraphAndTensors tensors) : tensors(std::move(tensors)) {} +struct CacheEntry { + explicit CacheEntry(GraphAndTensors graph_and_tensors) + : graph_and_tensors(std::move(graph_and_tensors)) {} - GraphAndTensors tensors; - std::once_flag plans_built; + GraphAndTensors graph_and_tensors; + std::once_flag build_plans_once; }; // One build site's cache. Process-wide rather than per-thread so a graph is reused across threads @@ -120,7 +96,7 @@ struct CachedGraph { template struct GraphCache { std::mutex mutex; // guards everything below - std::map>> entries; + std::map>> entries; }; // Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: @@ -128,22 +104,22 @@ struct GraphCache { // and both backends, so it is defined once here; `backend` and `pass` only name the build site the // stage timers attribute the calls to. // -// Reports by throwing, and the throw carries cuDNN's message alone. That message is what the -// is_supported_* helpers return as the reason a backend was refused, so a bool would discard the -// one thing a support probe exists to produce -- and NVTE_ERROR would wrap it in the file, line -// and advice of an internal failure, which a backend refused for a plain reason is not. One kind -// of throw for every failure; see support_verdict() for why that distinction is not drawn. +// Reports by throwing, and the throw carries cuDNN's message alone. That message is what +// support_verdict() returns as the reason a backend was refused, so a bool would discard the one +// thing a support probe exists to produce -- and NVTE_ERROR would wrap it in the file, line and +// advice of an internal failure, which a backend refused for a plain reason is not. One kind of +// throw for every failure; see support_verdict() for why that distinction is not drawn. // // graph.build_plans() and graph.execute() sit outside this function: they commit real resources, -// and the plan build belongs to whoever executes the graph, once. See CachedGraph. -inline void query_support(graph_cache_debug::Backend backend, Pass pass, - cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { +// and the plan build belongs to whoever executes the graph, once. See CacheEntry. +inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, + cudnnHandle_t handle) { auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { cudnn_frontend::error_t error; graph_cache_debug::record_time(backend, pass, stage, [&] { error = call(); }); if (error.is_good()) return; // cuDNN normally explains itself; fall back to the call's name so that a refusal can never - // arrive as an empty string, which the is_supported_* helpers would read as an endorsement. + // arrive as an empty string, which support_verdict() would read as an endorsement. throw std::runtime_error(error.err_msg.empty() ? std::string(call_name) + " failed." : error.err_msg); }; @@ -181,22 +157,22 @@ inline void query_support(graph_cache_debug::Backend backend, Pass pass, // lock cache.mutex // entries[key]? found -> copy the shared_ptr // unlock -// record_cache_lookup(HIT | MISS) +// record_hit_miss(HIT | MISS) // // HIT -> return the entry -// MISS -> build() outside the lock, so builds of unrelated -// query_support() keys proceed concurrently -// ok -> lock, insert, unlock, return the inserted entry, which on a lost race -// is the winner's +// MISS -> build(), record_create_graph() outside the lock, so builds of +// query_support() unrelated keys proceed concurrently +// ok -> record_cache_graph(), then lock, insert, unlock; return the inserted +// entry, which on a lost race is the winner's // throw -> propagates; nothing is stored, so the key is built again if it comes back template -std::shared_ptr> lookup_or_cache_graph( - GraphCache &cache, const FusedAttnConfig &key, - graph_cache_debug::Backend backend, Pass pass, cudnnHandle_t handle, BuildFn &&build) { - using Entry = CachedGraph; +std::shared_ptr> cache_graph(GraphCache &cache, + const FusedAttnConfig &key, + Backend backend, Pass pass, + cudnnHandle_t handle, BuildFn &&build) { using graph_cache_debug::LookupResult; - std::shared_ptr cached; + std::shared_ptr> cached; { std::lock_guard lock(cache.mutex); auto it = cache.entries.find(key); @@ -205,33 +181,111 @@ std::shared_ptr> lookup_or_cache_graph( // Recorded after the lock is dropped, so writing a trace line cannot hold up threads querying // other keys. The counters stay exact, but two lookups that raced can be recorded in the opposite // order, so a level-2 trace is the set of lookups that happened, not their sequence. - graph_cache_debug::record_cache_lookup( + graph_cache_debug::record_hit_miss( backend, pass, cached != nullptr ? LookupResult::Hit : LookupResult::Miss, key); if (cached != nullptr) return cached; - // A failure propagates with cuDNN's message and leaves nothing behind. It raised a MISS and no - // CREATE_GRAPH, which is what makes miss - create_graph the count of builds that ended this way. - auto entry = std::make_shared(build()); + // No backend refuses a configuration from in here any more -- TE's own FP8 rules moved to + // nvte_get_fused_attn_backend_v2, which answers them with a reason instead of throwing -- so + // `build` returning is now the ordinary case and miss == create_graph in a run that behaves. + // A throw that does get out still propagates with its message and leaves nothing behind, which + // is what keeps miss - create_graph worth printing: it should read zero. + auto entry = std::make_shared>(build()); + graph_cache_debug::record_create_graph(backend, pass); // Every site's tensor tuple leads with its graph, the one element this needs. A tuple ordered // otherwise would fail to compile rather than quietly validate the wrong object. - query_support(backend, pass, *std::get<0>(entry->tensors), handle); - graph_cache_debug::record_graph_created(backend, pass); - { - std::lock_guard lock(cache.mutex); - // On a losing race the insert does nothing: the shared_ptr this thread built is dropped with - // its graph, and what comes back is the winner's entry. - auto inserted = cache.entries.insert({key, std::move(entry)}); - return inserted.first->second; + // + // The two counters bracket this call deliberately: a graph cuDNN refuses throws here, having + // already recorded its CREATE_GRAPH and never reaching CACHE_GRAPH, so the gap between those two + // columns is cuDNN's refusals alone. + query_support(backend, pass, *std::get<0>(entry->graph_and_tensors), handle); + // Recorded on cuDNN's verdict rather than on the insert below, so the column counts the graphs + // cuDNN agreed to run. That is the question worth a counter; how many entries a map ended up + // holding is not, and tying it to the insert made a lost race -- which discards a supported graph + // and takes the winner's -- read as a miscount rather than as the duplicate work it is. + graph_cache_debug::record_cache_graph(backend, pass); + std::lock_guard lock(cache.mutex); + // On a losing race the insert does nothing: the shared_ptr this thread built is dropped with + // its graph, and what comes back is the winner's entry. + return cache.entries.insert({key, std::move(entry)}).first->second; +} + +// A backend's graph cache for one pass, and the only route to it. Both the execution path and the +// support probe come through here, so a probe leaves behind exactly the entry a later execution +// finds. That is what lets the probe's answer describe the graph that actually runs, rather than a +// separately built lookalike. +// +// The cache is this instantiation's static local, so the callers that name one triple share one cache, and each triple gets its own. Naming the triple is now what +// picks the cache, where before there was a per-backend function per pass to call. +// +// `kCreateGraphFn` is a template parameter rather than a `CreateFn &&` argument on purpose. As a +// parameter it makes the creator part of the instantiation, keeping the cache identified by the +// function that fills it. Passed as an argument, each distinct lambda type would instantiate its +// own copy of this function with its own static cache, and the two call sites for a pass would +// quietly stop sharing entries. +template +auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + static GraphCache cache; + // Asserted once here for both the key and the graph, which read the same derived fields. + check_derived(cfg); + return cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, + [&] { return kCreateGraphFn(cfg); }); +} + +// Whether cuDNN can run the graph this config asks for, in one direction: the empty string if it +// can, otherwise cuDNN's own account of why not, which the backend selector reports to the caller. +// This is the whole of what support_verdict_f16 and support_verdict_fp8 do; they exist only to +// reach their own translation unit's graph builders, which is also where a runtime direction turns +// into the compile-time one this needs. +// +// Named for what it returns rather than the question it answers: support is the empty string, so +// an is_supported() spelling would read backwards wherever the result is tested. +// +// The question is answered by building the graph, which is where every rejection comes from -- +// there is no separate list of rules to keep in step with the builder. The graph goes into the same +// cache the execution path reads, so the work is not thrown away and what was checked is what will +// run. It stops short of graph.build_plans(), the expensive step, which the first execution of the +// graph does instead; see CacheEntry. +// +// A refusal, by contrast, is not cached: nothing is stored for a key cuDNN rejected, so asking the +// same question again pays for the build again. See cache_graph. +// +// Refusals and failures on the way to a verdict read alike, because CUDNN_BACKEND_API_FAILED -- +// raised for any non-success cudnnStatus_t -- cannot separate CUDNN_STATUS_NOT_SUPPORTED from +// CUDNN_STATUS_ALLOC_FAILED. Either way this backend cannot serve this call, and either way what +// the caller wants is the message. +// +// The direction is named by the caller rather than read off the config: a config arriving from a +// framework has both check_for_*_support set, so it cannot say which graph is being probed. +template +std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + // Built only where it is used, on the two paths where a refusal arrived without a message of its + // own. Support is signalled by returning the empty string, so an empty refusal would otherwise + // read as an endorsement. + auto label = [] { + return std::string("support_verdict<") + graph_cache_debug::backend_name(kBackend) + ", " + + graph_cache_debug::pass_name(kPass) + ">"; + }; + try { + get_graph(cfg, handle); + return ""; + } catch (const std::exception &e) { + const char *reason = e.what(); + if (reason != nullptr && reason[0] != '\0') return reason; + return label() + ": rejected without a reason."; + } catch (...) { + return label() + ": unknown failure."; } } -// Runs graph.build_plans(), the plan build that lookup_or_cache_graph() left undone, once per -// entry. Named for the frontend call it wraps; the once-per-entry part is the whole reason it is a -// function rather than that call. +// Runs graph.build_plans(), the plan build that cache_graph() left undone, once per entry. Named +// for the frontend call it wraps; the once-per-entry part is the whole reason it is a function +// rather than that call. // // Call only when the graph is about to be executed, which is why this is a separate step rather // than the tail of the lookup: a support query builds entries nothing ever runs, and kernel -// compilation is the most expensive of the five frontend calls. See CachedGraph for why the flag +// compilation is the most expensive of the five frontend calls. See CacheEntry for why the flag // lives inside the entry and what a throw here leaves behind. // // Splitting the build in two means the thread that finishes it is often not the thread that started @@ -252,13 +306,12 @@ std::shared_ptr> lookup_or_cache_graph( // - graph.execute() is called with the running thread's own handle, so a handle is never used by // two threads at once, which is what cuDNN asks in return for letting them share a plan. template -void build_plans(graph_cache_debug::Backend backend, Pass pass, - CachedGraph &entry) { - std::call_once(entry.plans_built, [&] { - cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.tensors); +void build_plans(Backend backend, Pass pass, CacheEntry &entry) { + std::call_once(entry.build_plans_once, [&] { + cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.graph_and_tensors); graph_cache_debug::record_time(backend, pass, graph_cache_debug::BuildStage::BuildPlans, [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); - graph_cache_debug::record_plans_built(backend, pass); + graph_cache_debug::record_build_plans(backend, pass); }); } diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index c19a2902b0..b26181e2ec 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -12,8 +12,8 @@ // follows is what maintaining this file needs. // // level 1 (events) : one line per event that happens once per distinct cache key -// (CREATE_GRAPH, BUILD_PLANS), plus the exit summary block and -// its stage timings. Low volume by construction. +// (CREATE_GRAPH, CACHE_GRAPH, BUILD_PLANS), plus the exit summary +// block and its stage timings. Low volume by construction. // level 2 (trace) : adds a line per cache lookup (HIT/MISS, with the normalized // key) and per execution (EXEC). High volume, and it serializes // threads on the stderr lock, which the stage timings are then @@ -27,19 +27,30 @@ // // One level-1 training step, line prefixes and trailing columns elided: // -// tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, ... +// tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, cache_graph=0, ... +// tid=0 dev=0 | f16 fwd CACHE_GRAPH | hit=0, miss=1, create_graph=1, cache_graph=1, ... // ===== summary begin ===== -// tid=0 dev=0 | f16 fwd | hit=5, miss=1, create_graph=1, ... -// tid=1 dev=0 | f16 bwd | hit=4, build_plans=1, exec=1, ... -// tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, ... +// tid=0 dev=0 | f16 fwd | hit=5, miss=1, create_graph=1, cache_graph=1, ... +// tid=1 dev=0 | f16 bwd | hit=4, build_plans=1, execute=1, ... +// tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, cache_graph=1, ... // f16 fwd build_plans | calls=1 | time= 262.104 ms/call // ===== summary end ===== // +// Those first two lines are one graph, before and after cuDNN was asked to support it, which is why +// a CREATE_GRAPH with no CACHE_GRAPH following it is a refusal -- the one event pattern that says a +// build was paid for and thrown away. +// // Rows for a site a thread never reached are left out rather than zeroed, which is // why tid=1 has a backward row and no forward one: in a PyTorch step the forward and // the backward's support probe run on the main thread, and the backward itself on the // autograd thread, which finds the graph that probe left behind. That split is why // the build identities hold on the totals rows and not on any single thread's. +// +// Reading this file: the interface is the four names under "vocabulary" and the six recorders at +// the bottom, and that is everything the rest of the library touches. In between, in namespace +// detail, is what they are built out of, in the order an event travels through it -- the gate, the +// counters, the line, the exit summary. A question about what the output means is answered by the +// counter definitions in the middle; a question about what to call is answered by the bottom. // ============================================================================ #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ @@ -66,6 +77,46 @@ namespace transformer_engine { namespace fused_attn { namespace graph_cache_debug { +// ============================================================================ +// The vocabulary a call site needs: which build site an event came from, and which build stage or +// lookup outcome it is reporting. These four names and the recorders at the bottom of the file are +// the whole interface; everything between them is machinery, in namespace detail. +// +// Backend and Pass are fused_attn's own, from config_and_params.h, so that a recorder and the key +// it prints share one notion of a site; taking the pair rather than the "fwd"/"bwd" strings this +// used to also turns a mistake at a call site into a compile error. +// +// Every recorder names both halves, since the counters are per site -- adding f16's builds into +// fp8's column would leave a run that drove both unable to say which paid for what. +// ============================================================================ + +inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } +inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } + +// The frontend calls that make up a build, in the order they run. `kCount` must stay last: it +// sizes the timing table, and detail::kStageNames is indexed by these values when the summary +// prints, so the two must be kept in the same order. +enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; + +// What a lookup found: an entry, or nothing. +enum class LookupResult { Miss, Hit }; + +// ============================================================================ +// Machinery: the gate, the counters, the formatting and the exit summary. Nothing outside this +// file names any of it. +// +// Reading order below is the order an event travels: whether to record at all, which site it +// belongs to, the counters it moves, the line it prints, and finally the summary that reports the +// lot at exit. +// ============================================================================ +namespace detail { + +// ============================================================================ +// The gate: whether this process records anything, and how it names itself when it does. Every +// answer here is fixed for the life of the process and read out of an initialized-once static, so +// the check a disabled build pays at each call site is one load and one branch. +// ============================================================================ + // Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=events, // 2=trace). Single read at startup, cached; when unset every call site pays one // cached-flag check and nothing else. @@ -92,13 +143,18 @@ inline int launcher_rank() { return rank; } -// Whether this process emits diagnostics. Every rank writes to the same stderr, and under -// data/tensor parallelism they run identical shapes, so emitting from all of them multiplies the -// volume by the world size to say the same thing. Hence rank 0 only by default, overridable with -// the ":" suffix. Context parallelism is the case worth overriding for: the ranks run +// Diagnostics are on at level >= 1, and only for the ranks the ":" suffix selects. Every +// rank writes to the same stderr, and under data/tensor parallelism they run identical shapes, so +// emitting from all of them multiplies the volume by the world size to say the same thing. Hence +// rank 0 only by default. Context parallelism is the case worth overriding for: the ranks run // different subsets of the per-step regimes, so their build counts genuinely differ. -inline bool rank_selected() { - static const bool selected = [] { +// +// Both inputs are fixed for the life of the process, so the whole verdict is one initialized-once +// static -- which is all the check every call site makes reads, the per-lookup path included. +// Unselected ranks skip the counters too, so they pay nothing beyond it. +inline bool enabled() { + static const bool on = [] { + if (debug_level() < 1) return false; const int rank = launcher_rank(); if (rank < 0) return true; // sole process, nothing to filter const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); @@ -116,21 +172,14 @@ inline bool rank_selected() { } return false; }(); - return selected; -} - -// Diagnostics are on at level >= 1, and only for the selected ranks. Unselected ranks skip the -// counters too, so they pay nothing beyond this check. Cached in its own flag rather than -// recomputed from the two above, so that the check every call site makes -- the per-lookup path -// included -- reads one initialized-once static instead of two. Both inputs are fixed for the -// life of the process. -inline bool enabled() { - static const bool on = debug_level() >= 1 && rank_selected(); return on; } -// Per-lookup / per-exec trace lines are gated behind level >= 2. -inline bool trace_enabled() { return debug_level() >= 2; } +// The gate on the per-lookup and per-execution trace lines: everything enabled() asks for, level 2 +// on top of it. Named for that conjunction, and testing it rather than just the level, so the +// answer holds wherever it is asked -- level 2 alone is true on a rank that emits nothing, which +// would make this read as "trace" on every rank in the job. +inline bool enabled_with_trace() { return enabled() && debug_level() >= 2; } // Names the emitting rank, without which the ranks sharing one stderr would be indistinguishable. // A run whose launcher exports no rank is left untagged rather than falling back to a pid, an @@ -159,21 +208,9 @@ inline unsigned thread_seq_id() { inline void register_summary_once(); // ============================================================================ -// The build site an event came from: f16 or fp8, forward or backward. Every recorder names both -// halves, since the counters are per site -- adding f16's builds into fp8's column would leave a -// run that drove both unable to say which paid for what. A pair of enums rather than the -// "fwd"/"bwd" strings this used to take also turns a mistake at a call site into a compile error. -// -// Backend::F16 is the arbitrary-seqlen backend; the max512 one keeps no graph cache. Pass is -// fused_attn::Pass, from config_and_params.h, so that a recorder and the key it prints share one -// notion of direction. +// Indexing the build site an event came from: f16 or fp8, forward or backward. // ============================================================================ -enum class Backend { F16, FP8 }; - -inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } -inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } - // Backend major, pass minor, so that the two passes of one backend are adjacent -- which is how // the counter lines and the summary rows present them, one backend at a time. constexpr size_t kSiteCount = 4; @@ -184,56 +221,68 @@ inline constexpr size_t site_index(Backend b, Pass p) { // ============================================================================ // Cache event counters, one block per build site. Each name is both the event tag on the line // that records it and the column carrying its running total: -// - create_graph: a graph created and cached for a miss, only as far as check_support(). +// - create_graph: a graph constructed for a miss, counted before cuDNN is asked to support it and +// so regardless of what cuDNN goes on to say about it. +// - cache_graph: one of those graphs cleared check_support(), so this is the graphs cuDNN agreed +// to run. Counted on that verdict rather than on the insert that follows it, so it says what +// cuDNN accepted and not how many entries the map holds; the two differ only when a build race +// is lost and a supported graph is discarded for the winner's. // - build_plans: a cached graph finished with graph.build_plans(), the kernel compilation that -// create_graph deferred. At most one per create_graph, paid by that graph's first execution +// cache_graph deferred. At most one per cache_graph, paid by that graph's first execution // rather than by the probe that built it. -// - exec: a graph execution call with valid runtime tensors. -// - hit: a lookup answered from the cache. Need not lead to an exec -- it can be a backend +// - execute: a graph execution cuDNN accepted, counted once the enqueue returns. Not a completed +// execution: the work is asynchronous, so a device-side fault is not reflected here. +// - hit: a lookup answered from the cache. Need not lead to an execution -- it can be a backend // availability check, or the workspace-sizing call of nvte_fused_attn_fwd/bwd, which has no // runtime tensors to run with. // - miss: a lookup the cache did not answer; triggers a graph build. // // Identities, holding by construction, so a violation is a bug in the cache or in the counting // rather than something the workload did: -// - hit + miss = every lookup, one per entry into lookup_or_cache_graph, which makes it the -// denominator for everything below. -// - miss >= create_graph, the difference being builds that threw, whether cuDNN refused the graph -// or could not reach a verdict. Nothing is cached for those, so this is the only place a -// refusal shows up; its reason goes to the framework instead. -// - create_graph >= build_plans, the gap being graphs a probe built that nothing has run. -// - exec > 0 implies build_plans > 0, every site calling build_plans() ahead of the -// workspace-sizing return, itself ahead of record_exec. Read backwards: a workspace-sizing -// call pays build_plans and never exec. +// - hit + miss = every lookup, one per call to cache_graph(), which makes it the denominator for +// everything below. (The function, not the column of the same name -- the column counts the +// subset of those calls that ended in an entry.) +// - miss >= create_graph >= cache_graph, where each drop is a build that threw. +// create_graph - cache_graph is what cuDNN refused, or could not reach a verdict on; nothing is +// cached for it, so this gap is where a refusal shows up, and the reason itself goes to the +// framework. miss - create_graph would be a backend refusing a configuration from inside its +// own build, and no backend does that any more -- TE's own FP8 rules over bias, ALiBi and the +// recipes it does not serve moved to nvte_get_fused_attn_backend_v2, which answers them with a +// reason rather than throwing. So this gap should read zero, and the column stays as the thing +// that says so: a nonzero miss - create_graph means a build threw where none is expected to. +// - cache_graph >= build_plans, the gap being graphs a probe built that nothing has run. +// - execute > 0 implies build_plans > 0, every site calling build_plans() ahead of the +// workspace-sizing return, itself ahead of record_execute. Read backwards: a workspace-sizing +// call pays build_plans and never execute. // - Both build identities belong to the totals rows, not to one thread's: the thread that builds // a graph need not compile its plans, and a PyTorch step splits exactly that way. // - Per-thread rows sum column by column to "tid=all dev=all", and the per-backend rows of one // pass to that pass's all-backends row. -// - A lost build race disturbs none of the above -- the loser records its own miss and its own -// create_graph, and the once_flag still permits one build_plans -- but it does break reading -// create_graph as the number of graphs cached. +// - A lost build race disturbs none of the above: the loser records its own miss, create_graph +// and cache_graph, having built a graph cuDNN did agree to run, and the once_flag still permits +// one build_plans on the winner's entry. What it costs is a build, which two MISS lines on one +// key is the way to see. // - Stage timing calls fall along validate >= build_operation_graph >= create_execution_plans >= // check_support, each drop being the builds that ended at the stage before, which localizes // where cuDNN refuses rather than only how long refusing took. // - The build_plans timing row can show more calls than the build_plans column, the difference // being plan builds that threw: the timer records while unwinding, the counter only on return. // -// Signatures, workload-dependent, so read rather than asserted: -// - After warmup only hit and exec should move; a late create_graph means something varies per +// Signatures, workload-dependent, so read rather than asserted. What a column stalling says about +// who rejected a configuration is the user-facing half of this and lives in docs/envvars.rst; what +// follows is what is worth knowing on top of it: +// - After warmup only hit and execute should move; a late create_graph means something varies per // step that need not. -// - Several hits per exec is normal, since selection, workspace sizing and execution all look -// the same key up; what matters is that the ratio stays flat. -// - exec / create_graph is the amortization figure, and a lower bound at that, a lost race adding -// a build without a graph. Single digits after a long run means the cache is not earning its -// keep. -// - miss climbing while create_graph stays put is a configuration cuDNN keeps refusing, each query -// paying a discarded build. It also says this site never runs fused, making it the pair to read -// when attention is slower than expected and nothing raised an error. +// - Several hits per execution is normal, since selection, workspace sizing and execution all +// look the same key up; what matters is that the ratio stays flat. +// - execute / cache_graph is the amortization figure, and a lower bound at that, a lost race +// counting a supported graph the cache did not keep. Single digits after a long run means the +// cache is not earning its keep. // - miss climbing without settling means the key space is not closing, and since the cache is // unbounded, every distinct key is held for the life of the process. // - A build count that looks doubled on a multi-device process usually is not: device_id is part // of the key, so the same shape on two devices is two entries. Read the dev column. -// - Two MISS lines with the same key, create_graph above the number of distinct keys, is that lost +// - Two MISS lines with the same key, cache_graph above the number of distinct keys, is that lost // race: wasted work rather than a bug, worth chasing only if it repeats. // - A level-2 trace is the set of lookups, not their order, the line being written after the // cache lock is dropped. @@ -241,8 +290,9 @@ inline constexpr size_t site_index(Backend b, Pass p) { struct EventCounters { std::atomic create_graph{0}; + std::atomic cache_graph{0}; std::atomic build_plans{0}; - std::atomic exec{0}; + std::atomic execute{0}; std::atomic hit{0}; std::atomic miss{0}; }; @@ -258,15 +308,17 @@ inline EventCounters &counters(Backend b, Pass p) { // snapshot of a moving count by nature. struct CounterSnapshot { uint64_t create_graph = 0; + uint64_t cache_graph = 0; uint64_t build_plans = 0; - uint64_t exec = 0; + uint64_t execute = 0; uint64_t hit = 0; uint64_t miss = 0; CounterSnapshot &operator+=(const CounterSnapshot &other) { create_graph += other.create_graph; + cache_graph += other.cache_graph; build_plans += other.build_plans; - exec += other.exec; + execute += other.execute; hit += other.hit; miss += other.miss; return *this; @@ -274,19 +326,26 @@ struct CounterSnapshot { // Whether this block saw nothing at all, which is what lets the summary leave out the rows // for a backend the run never used rather than printing zeros for it. - bool empty() const { return (create_graph | build_plans | exec | hit | miss) == 0; } + bool empty() const { + return (create_graph | cache_graph | build_plans | execute | hit | miss) == 0; + } }; inline CounterSnapshot snapshot(const EventCounters &c) { CounterSnapshot s; s.create_graph = c.create_graph.load(std::memory_order_relaxed); + s.cache_graph = c.cache_graph.load(std::memory_order_relaxed); s.build_plans = c.build_plans.load(std::memory_order_relaxed); - s.exec = c.exec.load(std::memory_order_relaxed); + s.execute = c.execute.load(std::memory_order_relaxed); s.hit = c.hit.load(std::memory_order_relaxed); s.miss = c.miss.load(std::memory_order_relaxed); return s; } +// ============================================================================ +// The same counters again, per thread, and the registry the exit summary walks to find them. +// ============================================================================ + // Per-thread counters, one block per build site, so the summary can break every column down by // thread and backend: in the single-process context-parallel case each device is driven by its own // thread, and under PyTorch this separates the main thread from the autograd one. @@ -337,6 +396,33 @@ inline EventCounters &thread_counters(Backend b, Pass p) { return thread_counters().sites[site_index(b, p)]; } +// ============================================================================ +// Turning a counter block into a line, and getting a line out. One formatter, shared by the event +// lines and the summary rows, so that the two cannot drift into presenting the same columns +// differently, and one writer, so that everything here reaches stderr the same way. +// ============================================================================ + +// The one place diagnostics reach stderr, and the reason it exists: the first line this process +// writes carries a leading newline. Diagnostics share stderr with whatever the framework is +// printing, and a test runner's progress output has no trailing newline of its own, so without +// this the first line continues someone else's -- which on a level-2 trace line, long enough to +// wrap already, leaves no way to find where it starts. Where the previous output did end cleanly +// the prefix reads as a blank line setting the diagnostics apart from it. +// +// One fwrite per line either way: a rank's summary block is assembled whole precisely so that +// concurrently exiting ranks do not interleave, and the extra allocation buys the same for the one +// line that gets the prefix. +inline void write_stderr(const std::string &text) { + static std::atomic first_line{true}; + if (first_line.exchange(false, std::memory_order_relaxed)) { + const std::string first = "\n" + text; + std::fwrite(first.data(), 1, first.size(), stderr); + } else { + std::fwrite(text.data(), 1, text.size(), stderr); + } + std::fflush(stderr); +} + // Format one counter block -- one pass of one backend -- as one line. One pass rather than both // because a line carrying the forward and backward columns together ran past 300 characters and // wrapped in most terminals; the two passes are adjacent rows instead. @@ -358,9 +444,10 @@ inline std::string format_counter_line(const char *tid_field, const char *dev_fi char buf[512]; std::snprintf(buf, sizeof(buf), "[FUSED-ATTN-CACHE] %s%-7s %-7s | %s | hit=%4" PRIu64 ", miss=%4" PRIu64 - ", create_graph=%4" PRIu64 ", build_plans=%4" PRIu64 ", exec=%4" PRIu64 "\n", + ", create_graph=%4" PRIu64 ", cache_graph=%4" PRIu64 ", build_plans=%4" PRIu64 + ", execute=%4" PRIu64 "\n", rank_tag().c_str(), tid_field, dev_field, label, c.hit, c.miss, c.create_graph, - c.build_plans, c.exec); + c.cache_graph, c.build_plans, c.execute); return std::string(buf); } @@ -379,58 +466,40 @@ inline void print_counters(Backend b, Pass p, const char *event) { std::snprintf(label, sizeof(label), "%s %s %-12s", backend_name(b), pass_name(p), event); std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); - const std::string line = - format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p))); - std::fputs(line.c_str(), stderr); - std::fflush(stderr); -} - -// A graph created, taken through check_support() and cached. Call after that, from the miss path -// that did it -- after, because a graph cuDNN refuses throws instead of arriving here, which is -// what makes miss - create_graph the count of refused builds. -inline void record_graph_created(Backend b, Pass p) { - if (!enabled()) return; - register_summary_once(); - counters(b, p).create_graph.fetch_add(1, std::memory_order_relaxed); - thread_counters(b, p).create_graph.fetch_add(1, std::memory_order_relaxed); - print_counters(b, p, "CREATE_GRAPH"); + write_stderr(format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p)))); } -// The graph.build_plans() a create_graph deferred, now completed. Call from inside the -// std::call_once that runs it, and after the call returns rather than before: it throws without -// setting the once_flag, leaving a later execution to retry, so counting on the way out keeps this -// a count of graphs that reached a runnable state. -inline void record_plans_built(Backend b, Pass p) { - if (!enabled()) return; - register_summary_once(); - counters(b, p).build_plans.fetch_add(1, std::memory_order_relaxed); - thread_counters(b, p).build_plans.fetch_add(1, std::memory_order_relaxed); - print_counters(b, p, "BUILD_PLANS"); -} +// ============================================================================ +// What the recorders at the bottom of the file are made of: moving one column, and naming a +// lookup's outcome. +// ============================================================================ -inline void record_exec(Backend b, Pass p) { - if (!enabled()) return; +// The body every recorder shares: gate, register the exit summary, and add one to `column` +// in both the process-wide block and this thread's. Returns whether diagnostics are on at all, so +// that a caller can skip building a line nobody will read. +// +// Both blocks or neither. A recorder that moved one and not the other would leave the per-thread +// rows failing to add up to the totals row, which the summary presents as an invariant, and the +// discrepancy would look like a threading bug in the cache rather than a miscount here. +inline bool record_counter(Backend b, Pass p, std::atomic EventCounters::*column) { + if (!enabled()) return false; register_summary_once(); - counters(b, p).exec.fetch_add(1, std::memory_order_relaxed); - thread_counters(b, p).exec.fetch_add(1, std::memory_order_relaxed); - // The per-exec line fires on every execution; keep it out of the level-1 path. - if (!trace_enabled()) return; - print_counters(b, p, "EXEC"); + (counters(b, p).*column).fetch_add(1, std::memory_order_relaxed); + (thread_counters(b, p).*column).fetch_add(1, std::memory_order_relaxed); + return true; } -// What a lookup found: an entry, or nothing. -enum class LookupResult { Miss, Hit }; - -// The column a lookup lands in. Written as a switch with no default so that adding an outcome -// fails to compile here rather than being silently counted as a miss. -inline std::atomic &lookup_column(EventCounters &c, LookupResult result) { +// The column a lookup lands in, and the tag naming it. Both are written as a switch with no +// default so that adding an outcome fails to compile here rather than being silently counted as +// a miss. +inline std::atomic EventCounters::*lookup_column(LookupResult result) { switch (result) { case LookupResult::Hit: - return c.hit; + return &EventCounters::hit; case LookupResult::Miss: break; } - return c.miss; + return &EventCounters::miss; } inline const char *lookup_name(LookupResult result) { @@ -443,62 +512,6 @@ inline const char *lookup_name(LookupResult result) { return "MISS"; } -// `key` is the normalized cache key -- make_cache_key(pass)'s output, the exact value looked up -- -// not the execution config it came from. HIT/MISS is decided by comparing keys, so a trace of -// anything else cannot explain its own outcome: the pre-normalization config would show identical -// lines with opposite outcomes, and differing lines that both hit. Diffing two MISS lines here -// names exactly the fields responsible for the extra build. -// -// The cost is that overwritten fields are no longer visible in their original form: attn_scale -// reads 1, ragged num_tokens read 0, max_seqlen and batch_size read their bucketed values. -inline void record_cache_lookup(Backend b, Pass p, LookupResult result, - const FusedAttnConfig &key) { - if (!enabled()) return; - register_summary_once(); - lookup_column(counters(b, p), result).fetch_add(1, std::memory_order_relaxed); - lookup_column(thread_counters(b, p), result).fetch_add(1, std::memory_order_relaxed); - // The per-lookup config dump is the highest-volume line (one per cache lookup); - // keep it out of the level-1 path and off the stderr lock unless tracing. - if (!trace_enabled()) return; - std::fprintf( - stderr, - "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | 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", - rank_tag().c_str(), thread_seq_id(), key.device_id, backend_name(b), pass_name(p), - lookup_name(result), static_cast(key.is_training), static_cast(key.deterministic), - static_cast(key.cuda_graph), static_cast(key.return_max_logit), - static_cast(key.check_for_forward_support), static_cast(key.attn_mask_type), - static_cast(key.bias_type), static_cast(key.window_size_left), - static_cast(key.window_size_right), static_cast(key.bottom_right_diagonal), - static_cast(key.softmax_type), static_cast(key.scaling_mode), - static_cast(key.dropout), static_cast(key.attn_scale), - static_cast(key.qkv_dtype), static_cast(key.o_dtype), - static_cast(key.do_dtype), static_cast(key.dqkv_dtype), - static_cast(key.qkv_layout), static_cast(key.o_format), - static_cast(key.do_format), static_cast(key.dqkv_layout), - static_cast(key.qkv_scale_inv_format), static_cast(key.do_scale_inv_format), - static_cast(key.batch_size), static_cast(key.num_attn_heads), - static_cast(key.num_gqa_groups), static_cast(key.head_dim_qk), - static_cast(key.head_dim_v), static_cast(key.max_seqlen_q), - static_cast(key.max_seqlen_kv), static_cast(key.num_tokens_q), - static_cast(key.num_tokens_kv), static_cast(key.bucketed_batch_size), - static_cast(key.bucketed_num_tokens_q), - static_cast(key.bucketed_num_tokens_kv), static_cast(key.num_pages_k), - static_cast(key.num_pages_v), static_cast(key.page_size_k), - static_cast(key.page_size_v), static_cast(key.max_pages_per_seq_k), - static_cast(key.max_pages_per_seq_v), static_cast(key.bias_batch_size), - static_cast(key.bias_num_heads), static_cast(key.bias_seqlen_q), - static_cast(key.bias_seqlen_kv)); -} - // ============================================================================ // Graph build timings. // @@ -515,10 +528,8 @@ inline void record_cache_lookup(Backend b, Pass p, LookupResult result, // stage mean as where build time goes in aggregate, not as any one build's cost. // ============================================================================ -// The frontend calls that make up a build, in the order they run. `kCount` must -// stay last: it sizes the table below. `kStageNames` is indexed by these values -// when the summary prints, so the two must be kept in the same order. -enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; +// Indexed by BuildStage when the summary prints, so it must stay in that enum's order and carry +// one name per stage ahead of its kCount sentinel. inline constexpr const char *kStageNames[] = { "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; @@ -572,112 +583,232 @@ struct ScopedBuildTimer { } }; -// Record how long `fn` takes as `stage` of the given build site. Unlike the record_* helpers above -// this wraps the work rather than reporting on work already done, which is the point: preferred -// over a ScopedBuildTimer at the call site because the measured region is exactly the call passed -// in, so surrounding work cannot drift into it as that code changes. -template -inline void record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { - ScopedBuildTimer scoped(b, p, stage); - fn(); -} - // ============================================================================ // Summary: on process exit, print cache event counters and graph build timings. +// +// Each section below appends its rows to the block the handler is assembling, in the order they +// are printed: per-thread rows, then totals, then stage timings. Split into named pieces rather +// than written inline because they are read one at a time -- a question about the output is a +// question about one of these -- and because the registration itself is already three constructs +// deep (an initialized-once static holding an atexit handler) before any row logic joins it. // ============================================================================ + +// The two backends that keep a cache, in the order every part of the summary walks them. +inline constexpr Backend kSummaryBackends[] = {Backend::F16, Backend::FP8}; + +// Names one build site for a summary row. No padding: the site name is exactly the width of the +// column there, unlike the event lines, which pad it to keep their counters aligned. +inline std::string site_label(Backend b, Pass p) { + return std::string(backend_name(b)) + " " + pass_name(p); +} + +// How many backends the run actually drove. Decides whether the across-backend rows are worth +// printing: with one backend they would repeat that backend's own rows verbatim. +inline size_t active_backend_count() { + size_t active = 0; + for (const Backend b : kSummaryBackends) { + if (!snapshot(counters(b, Pass::Fwd)).empty() || !snapshot(counters(b, Pass::Bwd)).empty()) { + ++active; + } + } + return active; +} + +// Per-thread breakdown, sorted by tid, one row per build site that thread drove. Sites it never +// reached are left out, for the reason an unused backend is: a row of zeros says nothing. +inline void append_thread_rows(std::string &block) { + std::lock_guard lock(thread_registry_mutex()); + std::vector blocks = thread_registry(); + std::sort(blocks.begin(), blocks.end(), + [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); + for (const ThreadCounters *tc : blocks) { + char tid_field[16]; + char dev_field[16]; + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); + std::snprintf(dev_field, sizeof(dev_field), "dev=%d", + tc->device.load(std::memory_order_relaxed)); + for (const Backend b : kSummaryBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot c = snapshot(tc->sites[site_index(b, p)]); + if (c.empty()) continue; + block += format_counter_line(tid_field, dev_field, site_label(b, p).c_str(), c); + } + } + } +} + +// Totals, printed after the per-thread rows so they read as their sum: one row per build site, +// then one per pass across the backends when the run used more than one. Both come from the +// process-wide counters rather than by adding up the rows above, so the two agreeing is a check +// on the counting rather than an artifact of it. +inline void append_total_rows(std::string &block) { + CounterSnapshot all_fwd; + CounterSnapshot all_bwd; + for (const Backend b : kSummaryBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot c = snapshot(counters(b, p)); + (p == Pass::Fwd ? all_fwd : all_bwd) += c; + if (c.empty()) continue; + block += format_counter_line("tid=all", "dev=all", site_label(b, p).c_str(), c); + } + } + if (active_backend_count() <= 1) return; + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot &c = (p == Pass::Fwd ? all_fwd : all_bwd); + if (c.empty()) continue; + block += format_counter_line("tid=all", "dev=all", (std::string("all ") + pass_name(p)).c_str(), + c); + } +} + +// Mean time per call for each stage of each build site, skipping stages nothing reached. A mean is +// all the sums kept can support; see the section above for why that is the right figure to read. +inline void append_stage_rows(std::string &block) { + for (const Backend b : kSummaryBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + const StageTiming &t = stage_timing(b, p, static_cast(i)); + 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; + char line[288]; + std::snprintf( + line, sizeof(line), + "[FUSED-ATTN-CACHE] %s%-3s %-3s %-22s | calls=%" PRIu64 " | time=%9.3f ms/call\n", + rank_tag().c_str(), backend_name(b), pass_name(p), kStageNames[i], n, total_ms / n); + block += line; + } + } + } +} + inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; // Built in memory and emitted with one write, so that concurrently-exiting // processes (one per rank under torchrun) stay grouped rather than interleaving. - std::string block; - block += "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary begin =====\n"; - constexpr Backend kBackends[] = {Backend::F16, Backend::FP8}; - // A backend the run never reached is left out rather than reported as a row of zeros. - size_t active_backends = 0; - for (const Backend b : kBackends) { - if (!snapshot(counters(b, Pass::Fwd)).empty() || - !snapshot(counters(b, Pass::Bwd)).empty()) { - ++active_backends; - } - } - // Per-thread breakdown (sorted by tid), one row per build site that thread drove, with - // unreached sites left out for the same reason an unused backend is. - { - std::lock_guard lock(thread_registry_mutex()); - std::vector blocks = thread_registry(); - std::sort(blocks.begin(), blocks.end(), - [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); - for (const ThreadCounters *tc : blocks) { - char tid_field[16]; - char dev_field[16]; - std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); - std::snprintf(dev_field, sizeof(dev_field), "dev=%d", - tc->device.load(std::memory_order_relaxed)); - for (const Backend b : kBackends) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - const CounterSnapshot c = snapshot(tc->sites[site_index(b, p)]); - if (c.empty()) continue; - // No padding: a site name is exactly the width of the column on a summary row. - char label[32]; - std::snprintf(label, sizeof(label), "%s %s", backend_name(b), pass_name(p)); - block += format_counter_line(tid_field, dev_field, label, c); - } - } - } - } - // Totals last, so they read as the sum of the per-thread rows above: one row per build site, - // then a row per pass across the backends only when the run used more than one, since with a - // single backend those would repeat the rows above verbatim. - CounterSnapshot all_fwd; - CounterSnapshot all_bwd; - for (const Backend b : kBackends) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - const CounterSnapshot c = snapshot(counters(b, p)); - (p == Pass::Fwd ? all_fwd : all_bwd) += c; - if (c.empty()) continue; - char label[32]; - std::snprintf(label, sizeof(label), "%s %s", backend_name(b), pass_name(p)); - block += format_counter_line("tid=all", "dev=all", label, c); - } - } - if (active_backends > 1) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - const CounterSnapshot &c = (p == Pass::Fwd ? all_fwd : all_bwd); - if (c.empty()) continue; - char label[32]; - std::snprintf(label, sizeof(label), "all %s", pass_name(p)); - block += format_counter_line("tid=all", "dev=all", label, c); - } - } - for (const Backend b : kBackends) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { - const BuildStage s = static_cast(i); - const StageTiming &t = stage_timing(b, p, 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; - char line[288]; - std::snprintf( - line, sizeof(line), - "[FUSED-ATTN-CACHE] %s%-3s %-3s %-22s | calls=%" PRIu64 " | time=%9.3f ms/call\n", - rank_tag().c_str(), backend_name(b), pass_name(p), kStageNames[i], n, total_ms / n); - block += line; - } - } - } - block += "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary end =====\n"; - std::fwrite(block.data(), 1, block.size(), stderr); - std::fflush(stderr); + const std::string marker = "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary "; + std::string block = marker + "begin =====\n"; + append_thread_rows(block); + append_total_rows(block); + append_stage_rows(block); + block += marker + "end =====\n"; + write_stderr(block); }); return true; }(); (void)registered; } +} // namespace detail + +// ============================================================================ +// The recorders: everything a call site calls. Each takes the build site it is reporting for, adds +// one to that site's column, and prints a line when the level asks for it. +// +// Every one of them is called after the event it names, never before, so that a column counts what +// happened rather than what was attempted. That is what gives the gaps between columns their +// meaning: an event that can fail partway -- a build cuDNN refuses, an execution whose setup throws +// first -- leaves the earlier column moved and the later one not. +// +// record_time is the exception, and only because timing cannot be done after the fact: it brackets +// the call it measures, and accumulates while unwinding so that a stage which throws is still +// timed. Its timing rows can therefore outnumber the matching counter column. +// +// Which of them belongs where in the cache's flow is documented on each below and in graph_cache.h +// at the call sites. +// ============================================================================ + +// A graph constructed for a miss, whatever cuDNN goes on to make of it. Call from the miss path +// that built it, as soon as construction returns and before check_support() is asked. Before, +// because construction is where a backend would refuse a configuration on its own rules, and such +// a build never gets here -- which is what makes miss - create_graph builds that failed on this +// side of cuDNN. No backend does that now, so the gap is there to read as zero. +inline void record_create_graph(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::create_graph)) { + detail::print_counters(b, p, "CREATE_GRAPH"); + } +} + +// A created graph that cleared check_support(), so this counts the graphs cuDNN agreed to run. Call +// as soon as that verdict returns, ahead of the insert: a refused graph throws in between, leaving +// its CREATE_GRAPH unanswered, which is what makes create_graph - cache_graph cuDNN's refusals. +// Deliberately not the insert, so that a lost race reads as the extra build it is rather than as a +// count that disagrees with the size of the cache. +inline void record_cache_graph(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::cache_graph)) { + detail::print_counters(b, p, "CACHE_GRAPH"); + } +} + +// The graph.build_plans() a cache_graph deferred, now completed. Call from inside the +// std::call_once that runs it, and after the call returns rather than before: it throws without +// setting the once_flag, leaving a later execution to retry, so counting on the way out keeps this +// a count of graphs that reached a runnable state. +inline void record_build_plans(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::build_plans)) { + detail::print_counters(b, p, "BUILD_PLANS"); + } +} + +// An execution cuDNN accepted. Call after graph.execute() returns, as with the recorders above, so +// that a graph the surrounding setup never reached is not counted as having run -- the stream set +// and the cu_seqlens conversion kernels sit between the decision to execute and the execution, and +// either can throw. +// +// Accepted is as far as this can go. execute() enqueues on a stream and returns, so a fault the +// device raises later, surfacing at the next synchronization, still leaves the execution counted +// here. There is no synchronous completion point to hook without making the diagnostic change what +// it measures. +// +// Unlike the recorders above, this fires on every execution rather than once per distinct key, so +// its line is held back to level 2 while its column keeps counting. +inline void record_execute(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::execute) && + detail::enabled_with_trace()) { + detail::print_counters(b, p, "EXECUTE"); + } +} + +// `key` is the normalized cache key -- make_cache_key(pass)'s output, the exact value looked up -- +// not the execution config it came from. HIT/MISS is decided by comparing keys, so a trace of +// anything else cannot explain its own outcome: the pre-normalization config would show identical +// lines with opposite outcomes, and differing lines that both hit. Diffing two MISS lines here +// names exactly the fields responsible for the extra build. +// +// The cost is that overwritten fields are no longer visible in their original form: attn_scale +// reads 1, ragged num_tokens read 0, max_seqlen and batch_size read their bucketed values. +// +// This is the one line here not built from counters, so it does not go through +// format_counter_line: which fields it names is FusedAttnConfig::key_debug_string()'s to say, +// alongside the operator< that decides what a key compares on in the first place. +inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedAttnConfig &key) { + // The per-lookup config dump is the highest-volume line (one per cache lookup); + // keep it out of the level-1 path and off the stderr lock unless tracing. + if (!detail::record_counter(b, p, detail::lookup_column(result)) || + !detail::enabled_with_trace()) { + return; + } + char prefix[128]; + std::snprintf(prefix, sizeof(prefix), + "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | ", + detail::rank_tag().c_str(), detail::thread_seq_id(), key.device_id, backend_name(b), + pass_name(p), detail::lookup_name(result)); + detail::write_stderr(prefix + key.key_debug_string() + "\n"); +} + +// Record how long `fn` takes as `stage` of the given build site. Unlike the recorders above this +// wraps the work rather than reporting on work already done, which is the point: the measured +// region is exactly the call passed in, so surrounding work cannot drift into it as that code +// changes. Stage timings feed the summary only; they print no line of their own. +template +inline void record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { + detail::ScopedBuildTimer scoped(b, p, stage); + fn(); +} + } // namespace graph_cache_debug } // namespace fused_attn } // namespace transformer_engine From fbc09d4682c5aa8f143a8c297d8a8fc1e4121167 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 20 Aug 2026 15:52:50 +0000 Subject: [PATCH 84/88] [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 | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index b26181e2ec..7770dbad07 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -656,8 +656,8 @@ inline void append_total_rows(std::string &block) { for (const Pass p : {Pass::Fwd, Pass::Bwd}) { const CounterSnapshot &c = (p == Pass::Fwd ? all_fwd : all_bwd); if (c.empty()) continue; - block += format_counter_line("tid=all", "dev=all", (std::string("all ") + pass_name(p)).c_str(), - c); + block += + format_counter_line("tid=all", "dev=all", (std::string("all ") + pass_name(p)).c_str(), c); } } From 596c22729cb850165209f756d8b5dcaaa54f4a21 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:57:39 -0700 Subject: [PATCH 85/88] fix compile warnings and misc changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- docs/envvars.rst | 2 ++ tests/pytorch/attention/test_attention.py | 11 ++++++++--- .../common/fused_attn/fused_attn_fp8.cu | 10 +++------- transformer_engine/common/fused_attn/graph_cache.h | 4 ++-- .../common/fused_attn/graph_cache_debug.h | 11 ++++++++--- 5 files changed, 23 insertions(+), 15 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 85881f512c..4490320030 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -210,6 +210,8 @@ backend-selection overview. ``2`` additionally emits a per-lookup ``HIT``/``MISS`` line carrying the full cache key, and a per-execution ``EXECUTE`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. + Each line is written after the cache lock is released rather than under it, so that no thread waits on stderr while holding the cache. With several threads active this means the lines can appear in a different order than the lookups they report; the counters each line carries still increase in event order, and lines from a single thread are still in that thread's order. + By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. Has negligible overhead when unset. diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 8bd7235446..0202c1899d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -322,7 +322,7 @@ def test_dpa_checkpoint(dtype, model_configs, model): _CACHE_EVENT = re.compile( r"\[FUSED-ATTN-CACHE\]\s+(?:rank=\d+\s+\|\s+)?tid=\d+\s+dev=-?\d+\s+\|\s+" r"(?Pf16|fp8)\s+(?Pfwd|bwd)\s+" - r"(?PCREATE_GRAPH|BUILD_PLANS|EXEC|MISS|HIT)\b(?P.*)" + r"(?PCREATE_GRAPH|CACHE_GRAPH|BUILD_PLANS|EXECUTE|MISS|HIT)\b(?P.*)" ) _CACHE_PHASE = re.compile(r"\[CACHE-TEST\] phase=(?P\w+)") @@ -412,6 +412,9 @@ def count(phase, event, pass_name=pass_name): assert count("query", "MISS") == 1, f"{pass_name}: expected one cold miss{context}" assert count("query", "CREATE_GRAPH") == 1, f"{pass_name}: expected one build{context}" assert count("query", "BUILD_PLANS") == 0, f"{pass_name}: query compiled kernels{context}" + # The graph cuDNN took, which is the one the execution phases below go on to find. A + # build refused by check_support() would show up here as CREATE_GRAPH without this. + assert count("query", "CACHE_GRAPH") == 1, f"{pass_name}: build was not cached{context}" # Asking the identical question again must cost nothing. assert count("requery", "MISS") == 0, f"{pass_name}: repeated query missed{context}" @@ -429,7 +432,7 @@ def count(phase, event, pass_name=pass_name): assert ( count("exec", "CREATE_GRAPH") == 0 ), f"{pass_name}: execution rebuilt the graph{context}" - assert count("exec", "EXEC") >= 1, f"{pass_name}: fused attention never ran{context}" + assert count("exec", "EXECUTE") >= 1, f"{pass_name}: fused attention never ran{context}" assert count("exec", "BUILD_PLANS") == 1, f"{pass_name}: expected one plan build{context}" # softmax_scale reaches the graph as a pointer, not as a shape, so the key drops it: @@ -439,7 +442,9 @@ def count(phase, event, pass_name=pass_name): count("rescale", "CREATE_GRAPH") == 0 ), f"{pass_name}: attn_scale forced a build{context}" assert count("rescale", "BUILD_PLANS") == 0, f"{pass_name}: attn_scale recompiled{context}" - assert count("rescale", "EXEC") >= 1, f"{pass_name}: rescaled run did not execute{context}" + assert ( + count("rescale", "EXECUTE") >= 1 + ), f"{pass_name}: rescaled run did not execute{context}" # max_seqlen is a dimension the graph is built at, so it must miss -- once, for one # new graph, rather than invalidating what is already cached. diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 1c402f2676..91361f1d0a 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -84,7 +84,6 @@ static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { const int64_t d_v = static_cast(cfg.head_dim_v); const int64_t window_size_left = cfg.window_size_left; const int64_t window_size_right = cfg.window_size_right; - const bool is_training = cfg.is_training; const float dropout_probability = cfg.dropout; const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; const NVTE_QKV_Format o_format = cfg.o_format; @@ -367,7 +366,6 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de const int64_t b = static_cast(cfg.batch_size); // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; const bool is_dropout = cfg.is_dropout; const bool is_softmax_offset = cfg.is_softmax_offset; @@ -417,7 +415,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de variant_pack[amax_o] = devPtrAmaxO; } - /* if (is_bias) { + /* if (cfg.is_bias) { variant_pack[bias] = devPtrBias; } */ @@ -950,10 +948,8 @@ void fused_attn_fp8_bwd_impl( const bool is_O_in_F16 = !cfg.o_is_fp8; const int64_t b = static_cast(cfg.batch_size); - const int64_t h = static_cast(cfg.num_attn_heads); // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; - const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; const bool is_dropout = cfg.is_dropout; const bool is_softmax_offset = cfg.is_softmax_offset; @@ -1026,9 +1022,9 @@ void fused_attn_fp8_bwd_impl( variant_pack[descale_dO_t] = devPtrDescaledO_t; } - /* if (is_bias) { + /* if (cfg.is_bias) { variant_pack[bias] = devPtrBias; - if ((bias_b == 1) && (bias_h == h)) { + if ((bias_b == 1) && (bias_h == cfg.num_attn_heads)) { variant_pack[dBias] = devPtrdBias; } else { variant_pack[dBias] = nullptr; diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index 9e7b344342..521dcf6748 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -115,8 +115,8 @@ struct GraphCache { inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { - cudnn_frontend::error_t error; - graph_cache_debug::record_time(backend, pass, stage, [&] { error = call(); }); + const cudnn_frontend::error_t error = + graph_cache_debug::record_time(backend, pass, stage, [&] { return call(); }); if (error.is_good()) return; // cuDNN normally explains itself; fall back to the call's name so that a refusal can never // arrive as an empty string, which support_verdict() would read as an endorsement. diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 7770dbad07..882ec3b22a 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -15,7 +15,7 @@ // (CREATE_GRAPH, CACHE_GRAPH, BUILD_PLANS), plus the exit summary // block and its stage timings. Low volume by construction. // level 2 (trace) : adds a line per cache lookup (HIT/MISS, with the normalized -// key) and per execution (EXEC). High volume, and it serializes +// key) and per execution (EXECUTE). High volume, and it serializes // threads on the stderr lock, which the stage timings are then // measured under -- no timed region writes to stderr, so they // stay sound, but they read a little high. @@ -803,10 +803,15 @@ inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedA // wraps the work rather than reporting on work already done, which is the point: the measured // region is exactly the call passed in, so surrounding work cannot drift into it as that code // changes. Stage timings feed the summary only; they print no line of their own. +// +// Passes `fn`'s result back out so that a timed call reporting a value can be written as the +// initializer of that value. cuDNN's error_t is [[nodiscard]], and the alternative -- declaring the +// variable above the timing and assigning to it inside a void `fn` -- discards the assignment's own +// result, which the compiler counts as ignoring a nodiscard value. template -inline void record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { +inline decltype(auto) record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { detail::ScopedBuildTimer scoped(b, p, stage); - fn(); + return fn(); } } // namespace graph_cache_debug From 9609324f90c940e6c2738f2246398e5e79d548ed Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:14:02 -0700 Subject: [PATCH 86/88] WIP: trim code/comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- tests/pytorch/test_torch_compile.py | 1 - .../common/fused_attn/config_and_params.cpp | 403 +++++++--------- .../common/fused_attn/config_and_params.h | 322 ++++++------- .../common/fused_attn/fused_attn.cpp | 249 +++------- .../fused_attn_f16_arbitrary_seqlen.cu | 4 +- .../common/fused_attn/fused_attn_fp8.cu | 74 +-- .../common/fused_attn/graph_cache.h | 305 +++++------- .../common/fused_attn/graph_cache_debug.h | 447 +++++++----------- .../include/transformer_engine/fused_attn.h | 305 ++++++------ .../jax/cpp_extensions/attention.py | 125 +++-- transformer_engine/jax/csrc/extensions.h | 14 +- .../jax/csrc/extensions/attention.cpp | 133 ++++-- .../jax/csrc/extensions/pybind.cpp | 3 +- .../dot_product_attention/backends.py | 13 +- .../attention/dot_product_attention/utils.py | 8 +- .../pytorch/cpp_extensions/fused_attn.py | 13 +- .../pytorch/csrc/extensions/attention.cpp | 2 +- 17 files changed, 1083 insertions(+), 1338 deletions(-) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index e483c9188b..ca8f2b63fd 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -48,7 +48,6 @@ from transformer_engine.pytorch.attention.dot_product_attention.backends import ( UnfusedDotProductAttention, ) -from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index 8a28ba93ae..751f075a47 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -36,13 +36,17 @@ namespace fused_attn { // Forward declarations size_t get_max_batch_size(size_t batch_size); size_t get_max_tokens(size_t num_tokens); +DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads, + int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv, + int64_t head_dim_qk, int64_t head_dim_v); 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); - // convenience fields + // Convenience fields + qkv_format = nvte_get_qkv_format(qkv_layout); 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); @@ -63,20 +67,7 @@ void FusedAttnConfig::derive() { is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING); is_dropout = is_training && dropout != 0.0f; - // Both layouts describe variable-length sequences inside padded dimensions, so the mask is the - // only thing that tells cuDNN where the real tokens end; without it the graph attends to - // padding. Asserted here so that all four graph builders inherit the rule, and stated as a - // rejection rule in nvte_get_fused_attn_backend_v2 so that a support query answers rather than - // throws. Not conditioned on the cuDNN version: the requirement comes from what the dimensions - // mean, not from what any particular cuDNN can run. - if (is_paged_kv) { - NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); - } - if (is_ragged_q || is_ragged_kv) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - } - - // bucket the THD (ragged) batch and token counts + // 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); bucketed_batch_size = @@ -84,66 +75,63 @@ void FusedAttnConfig::derive() { 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, once per backend. Newer cuDNN SDPA can take sequence - // lengths directly as a cumulative tensor, which saves one kernel call; the frontend gates that - // on min(compile-time, runtime) cuDNN, so both versions are tested. The FP8 path needs newer - // versions of both than the F16 path, which is the whole reason there are two answers here. + // Use cu_seqlens vs actual_seqlens const size_t cudnn_runtime_version = cudnnGetVersion(); uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12500 && (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && !is_dropout; - // Frontend 1.26 supports fp8+cu_seqlens for the C++ API; the Python API needs 1.27. The dropout - // exclusion is not the F16 one restated: the frontend cannot combine dropout with stats - // generation in the fprop unified engine, so such a request would be routed to the old composite - // SDPA engine, which has no cu_seqlens support at all. Remove that term when it can be. fp8_uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12600 && (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && !is_dropout; - // What each pass stores, classified for the FP8 backend; see the fields for what reads them. - o_is_fp8 = (o_dtype == kNVTEFloat8E4M3 || o_dtype == kNVTEFloat8E5M2); - dqkv_is_fp8 = (dqkv_dtype == kNVTEFloat8E4M3 || dqkv_dtype == kNVTEFloat8E5M2); - - // packed vs dense dimensions for a ragged (THD) graph; SM8x and SM120 require dense, - // BHSD-like dimensions for the Stats/LSE auxiliary tensors and so take the dense path + is_o_in_fp8 = (o_dtype == kNVTEFloat8E4M3 || o_dtype == kNVTEFloat8E5M2); + is_dqkv_in_fp8 = (dqkv_dtype == kNVTEFloat8E4M3 || dqkv_dtype == kNVTEFloat8E5M2); + const bool is_o_in_f16 = (o_dtype == kNVTEFloat16 || o_dtype == kNVTEBFloat16); + const bool is_dqkv_in_f16 = (dqkv_dtype == kNVTEFloat16 || dqkv_dtype == kNVTEBFloat16); + + // Determine the FP8 recipe + is_tensor_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING); + is_delayed_scaling_fwd = is_tensor_scaling && is_o_in_fp8; + is_current_scaling_fwd = is_tensor_scaling && is_o_in_f16; + is_delayed_scaling_bwd = is_tensor_scaling && is_dqkv_in_fp8; + is_current_scaling_bwd = is_tensor_scaling && is_dqkv_in_f16; + is_mxfp8_fwd = is_mxfp8 && is_o_in_f16; + is_mxfp8_bwd = is_mxfp8 && is_dqkv_in_f16; + + // Whether packed graphs exist for THD const int sm_arch = cuda::sm_arch(cuda::current_device()); uses_packed_ragged_graph = cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; uses_ragged_stats = is_ragged_q && uses_packed_ragged_graph; - // sequence lengths the graph is built at + // Sequence lengths the graph is built at graph_max_seqlen_q = (is_ragged_q && uses_packed_ragged_graph) ? bucketed_num_tokens_q : max_seqlen_q; graph_max_seqlen_kv = (is_ragged_kv && uses_packed_ragged_graph) ? bucketed_num_tokens_kv : max_seqlen_kv; - // Batch size and ragged-offset width the graph is built at, for each direction. One condition - // decides all four: whether cuDNN is handed the caller's cu_seqlens* buffers untouched, which - // only the forward graph ever is. When it is, those buffers are what the graph has to match -- - // their [batch_size + 1] length, which a bucketed batch would read past the end of, and their - // int32 width. The backward graph always reads seqlens converted into our own workspace, so - // neither constraint reaches it and its two answers are the unconditional ones. Otherwise the - // batch is the bucketed one wherever a ragged layout is packed, so that one graph serves every - // batch in its bucket -- the same reason graph_max_seqlen_* stands in for the sequence lengths -- - // and the offset width is whichever the runtime supports, which is what lets older cuDNN - // runtimes work rather than fail. - // - // Kept as four adjacent assignments off shared locals because the pairs have to agree: a forward - // graph built at a bucketed batch while expecting offsets at the other width is the failure this - // arrangement exists to make visible. + // Ragged-offset width that this config needs + needs_64bit_ragged_offset = + (is_ragged_q || is_ragged_kv) && + fused_attn::get_ragged_offset_dtype( + layout_group, 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)) == DType::kInt64; const DType wide_ragged_offsets = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; - const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_packed_ragged_graph; ragged_offset_type_fwd = uses_cu_seqlens_directly ? DType::kInt32 : wide_ragged_offsets; ragged_offset_type_bwd = wide_ragged_offsets; + + // Batch size the graph is built at + const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_packed_ragged_graph; graph_batch_size_fwd = (buckets_the_batch && !uses_cu_seqlens_directly) ? bucketed_batch_size : batch_size; graph_batch_size_bwd = buckets_the_batch ? bucketed_batch_size : batch_size; - // elements per token for each ragged tensor + // Elements per token for each ragged tensor ragged_offset_mults = RaggedOffsetMultipliers( layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), static_cast(head_dim_qk), static_cast(head_dim_v)); - // paged KV dimensions + // Paged KV dimensions if (is_paged_kv) { if (num_pages_k == 0) { num_pages_k = static_cast(b); @@ -169,10 +157,7 @@ void FusedAttnConfig::derive() { } FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { - // Requires a derived config: every normalization below reads a derived field -- is_padding and - // is_causal_bottom_right, the is_ragged_* pair, the graph_max_seqlen_* dimensions, and the - // uses_* flags. A precondition rather than an assert, since every caller reaches this through - // get_graph(), which asserts it once for both the key and the graph. + check_derived(); FusedAttnConfig cache_cfg = *this; // Key the device ID for multi-GPU single-process runs @@ -187,18 +172,11 @@ FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { cache_cfg.bottom_right_diagonal = false; } - // Name the sequence lengths the graph is built at rather than the ones the caller asked about, - // so that every shape falling in the same bucket lands on the same entry. The two are equal - // unless a ragged layout is packed, which is why this is unconditional. Stated after the - // bottom_right_diagonal rule above, which is about the real geometry of the attention mask and - // would read bucketed token counts as sequence lengths if it ran after the substitution. + // Normalize sequence lengths the graph is built at cache_cfg.max_seqlen_q = cache_cfg.graph_max_seqlen_q; cache_cfg.max_seqlen_kv = cache_cfg.graph_max_seqlen_kv; - // Name the batch size the graph is built at, and drop the token counts the bucketing replaced. - // Reading the batch derive() recorded for this pass, rather than restating the rule that set it, - // is what keeps the key from naming a batch the graph was not built with -- the two directions - // bucket differently, and the graph builders read the same field for the same pass. + // Normalize batch size the graph is built at, and drop the token counts the bucketing replaced. if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_packed_ragged_graph) { cache_cfg.num_tokens_q = 0; cache_cfg.num_tokens_kv = 0; @@ -210,16 +188,10 @@ FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { cache_cfg.attn_scale = 1.0f; // cuda_graph never reaches a graph builder. Its one use is the cuDNN <= 9.15 rejection in - // nvte_get_fused_attn_backend_v2(), which is decided before the cache is consulted, so a - // configuration that gets this far builds the same graph either way. Left in the key it would - // give a workload that both captures and runs eagerly two entries for every configuration. + // nvte_get_fused_attn_backend_v2(). cache_cfg.cuda_graph = false; - // Restrict this direction's key to the fields its graph actually consumes, so no redundant - // graphs are built and no cache misses either. Keyed on the pass rather than on the - // check_for_*_support flags, so that a caller asking about both directions -- which every - // backend query from a framework does -- still gets a key each pass can find its own graph - // under, instead of one narrowed for neither. + // Normalize the fields its graph actually consumes if (pass == Pass::Fwd) { cache_cfg.do_dtype = kNVTEBFloat16; cache_cfg.dqkv_dtype = kNVTEBFloat16; @@ -231,22 +203,14 @@ FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { cache_cfg.return_max_logit = false; } - // The two flags say which directions the caller wanted probed, which the graph this key names - // does not depend on. Normalized so that a key is the same whether it came from a probe or from - // execution, and so that a level-2 trace line cannot claim a direction the key is not for. - cache_cfg.check_for_forward_support = pass == Pass::Fwd; - cache_cfg.check_for_backward_support = pass == Pass::Bwd; - return cache_cfg; } -std::string FusedAttnConfig::key_debug_string() const { - // Enums and sizes are printed as int64_t rather than by name, since the point is diffing two - // lines rather than reading one, and a numeric field cannot drift from a names table. +std::string FusedAttnConfig::to_string() const { char buf[1024]; std::snprintf( buf, sizeof(buf), - "train=%d det=%d cg=%d maxlogit=%d fwd=%d mask=%" PRId64 " bias=%" PRId64 " wl=%" PRId64 + "train=%d det=%d cg=%d maxlogit=%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 @@ -256,28 +220,27 @@ std::string FusedAttnConfig::key_debug_string() const { " npv=%" PRId64 " psk=%" PRId64 " psv=%" PRId64 " mppk=%" PRId64 " mppv=%" PRId64 " bias_b=%" PRId64 " bias_h=%" PRId64 " bias_sq=%" PRId64 " bias_skv=%" PRId64, static_cast(is_training), static_cast(deterministic), static_cast(cuda_graph), - static_cast(return_max_logit), static_cast(check_for_forward_support), - static_cast(attn_mask_type), static_cast(bias_type), - static_cast(window_size_left), static_cast(window_size_right), - static_cast(bottom_right_diagonal), static_cast(softmax_type), - static_cast(scaling_mode), static_cast(dropout), - static_cast(attn_scale), static_cast(qkv_dtype), - static_cast(o_dtype), static_cast(do_dtype), - static_cast(dqkv_dtype), static_cast(qkv_layout), - static_cast(o_format), static_cast(do_format), - static_cast(dqkv_layout), static_cast(qkv_scale_inv_format), - static_cast(do_scale_inv_format), static_cast(batch_size), - static_cast(num_attn_heads), static_cast(num_gqa_groups), - static_cast(head_dim_qk), static_cast(head_dim_v), - static_cast(max_seqlen_q), static_cast(max_seqlen_kv), - static_cast(num_tokens_q), static_cast(num_tokens_kv), - static_cast(bucketed_batch_size), static_cast(bucketed_num_tokens_q), - static_cast(bucketed_num_tokens_kv), static_cast(num_pages_k), - static_cast(num_pages_v), static_cast(page_size_k), - static_cast(page_size_v), static_cast(max_pages_per_seq_k), - static_cast(max_pages_per_seq_v), static_cast(bias_batch_size), - static_cast(bias_num_heads), static_cast(bias_seqlen_q), - static_cast(bias_seqlen_kv)); + static_cast(return_max_logit), static_cast(attn_mask_type), + static_cast(bias_type), static_cast(window_size_left), + static_cast(window_size_right), static_cast(bottom_right_diagonal), + static_cast(softmax_type), static_cast(scaling_mode), + static_cast(dropout), static_cast(attn_scale), + static_cast(qkv_dtype), static_cast(o_dtype), + static_cast(do_dtype), static_cast(dqkv_dtype), + static_cast(qkv_layout), static_cast(o_format), + static_cast(do_format), static_cast(dqkv_layout), + static_cast(qkv_scale_inv_format), static_cast(do_scale_inv_format), + static_cast(batch_size), static_cast(num_attn_heads), + static_cast(num_gqa_groups), static_cast(head_dim_qk), + static_cast(head_dim_v), static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), static_cast(num_tokens_q), + static_cast(num_tokens_kv), static_cast(bucketed_batch_size), + static_cast(bucketed_num_tokens_q), static_cast(bucketed_num_tokens_kv), + static_cast(num_pages_k), static_cast(num_pages_v), + static_cast(page_size_k), static_cast(page_size_v), + static_cast(max_pages_per_seq_k), static_cast(max_pages_per_seq_v), + static_cast(bias_batch_size), static_cast(bias_num_heads), + static_cast(bias_seqlen_q), static_cast(bias_seqlen_kv)); return std::string(buf); } @@ -823,6 +786,15 @@ void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, case kNVTEFusedAttnFwdParamsSoftmaxOffset: std::memcpy(buf, &p.SoftmaxOffset, 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 kNVTEFusedAttnFwdParamsCuSeqlensQ: std::memcpy(buf, &p.cu_seqlens_q, attr_size); break; @@ -844,30 +816,42 @@ void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, 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); + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); break; - case kNVTEFusedAttnFwdParamsAuxCtxTensors: - std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); break; case kNVTEFusedAttnFwdParamsIsTraining: bool_to_uint8(p.is_training, buf); break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + bool_to_uint8(p.return_max_logit, buf); + break; case kNVTEFusedAttnFwdParamsCudaGraph: bool_to_uint8(p.cuda_graph, buf); break; - case kNVTEFusedAttnFwdParamsReturnMaxLogit: - bool_to_uint8(p.return_max_logit, buf); + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); break; - case kNVTEFusedAttnFwdParamsAttnMaskType: - std::memcpy(buf, &p.attn_mask_type, attr_size); + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(buf, &p.dropout, 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; @@ -880,27 +864,6 @@ void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, 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; @@ -940,6 +903,15 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, case kNVTEFusedAttnFwdParamsSoftmaxOffset: std::memcpy(&p.SoftmaxOffset, 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 kNVTEFusedAttnFwdParamsCuSeqlensQ: std::memcpy(&p.cu_seqlens_q, buf, attr_size); break; @@ -961,30 +933,42 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, 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); + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); break; - case kNVTEFusedAttnFwdParamsAuxCtxTensors: - std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); break; case kNVTEFusedAttnFwdParamsIsTraining: uint8_to_bool(buf, p.is_training); break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + uint8_to_bool(buf, p.return_max_logit); + break; case kNVTEFusedAttnFwdParamsCudaGraph: uint8_to_bool(buf, p.cuda_graph); break; - case kNVTEFusedAttnFwdParamsReturnMaxLogit: - uint8_to_bool(buf, p.return_max_logit); + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); break; - case kNVTEFusedAttnFwdParamsAttnMaskType: - std::memcpy(&p.attn_mask_type, buf, attr_size); + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(&p.dropout, 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; @@ -997,27 +981,6 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, 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; @@ -1106,36 +1069,18 @@ void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); break; - case kNVTEFusedAttnBwdParamsCudaGraph: - bool_to_uint8(p.cuda_graph, buf); - break; - 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); + 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; - case kNVTEFusedAttnBwdParamsBottomRightDiagonal: - bool_to_uint8(p.bottom_right_diagonal, buf); + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); 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); break; @@ -1154,11 +1099,29 @@ void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsDOScaleInvFormat: std::memcpy(buf, &p.do_scale_inv_format, attr_size); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenQ: - std::memcpy(buf, &p.max_seqlen_q, attr_size); + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenKV: - std::memcpy(buf, &p.max_seqlen_kv, attr_size); + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_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 kNVTEFusedAttnBwdParamsDeterministic: + bool_to_uint8(p.deterministic, buf); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); break; case kNVTEFusedAttnBwdParamsWorkspace: std::memcpy(buf, &p.workspace, attr_size); @@ -1235,36 +1198,18 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); break; - case kNVTEFusedAttnBwdParamsCudaGraph: - uint8_to_bool(buf, p.cuda_graph); - break; - 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); + 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; - case kNVTEFusedAttnBwdParamsBottomRightDiagonal: - uint8_to_bool(buf, p.bottom_right_diagonal); + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); 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); break; @@ -1283,11 +1228,29 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, case kNVTEFusedAttnBwdParamsDOScaleInvFormat: std::memcpy(&p.do_scale_inv_format, buf, attr_size); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenQ: - std::memcpy(&p.max_seqlen_q, buf, attr_size); + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); break; - case kNVTEFusedAttnBwdParamsMaxSeqlenKV: - std::memcpy(&p.max_seqlen_kv, buf, attr_size); + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_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 kNVTEFusedAttnBwdParamsDeterministic: + uint8_to_bool(buf, p.deterministic); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); break; case kNVTEFusedAttnBwdParamsWorkspace: std::memcpy(&p.workspace, buf, attr_size); diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h index 5696e1e23b..7bfe3dc1ad 100644 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -21,22 +21,7 @@ namespace transformer_engine { namespace fused_attn { -// The pair that names one build site: whose graphs, and which of the two a config is being turned -// into. A site keeps its own graph cache and its own counters, so both halves travel together. -// -// Declared here, with the config, rather than with the cache or its diagnostics: they are the -// vocabulary those two share, and the config is where their consumers start -- make_cache_key() -// below turns a config into one site's key, and derive() fills the dimensions built from it. -// -// Backend::F16 is the arbitrary-seqlen backend; the max512 one keeps no graph cache, so it has no -// site here. Narrower than the public NVTE_Fused_Attn_Backend, and not a substitute for it: this -// names only the backends that build graphs. enum class Backend { F16, FP8 }; - -// Passed in rather than derived, because a config cannot say which graph is being built from it. -// check_for_forward_support and check_for_backward_support state which directions a caller wants -// probed, and a backend query arriving from a framework has both set, so they answer a different -// question -- see the comment on them below. enum class Pass { Fwd, Bwd }; struct FusedAttnConfig { @@ -94,77 +79,60 @@ 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 + // Internal fields: keyed + // + // device ID is not part of attribute serialization, i.e. internal, but it participates in + // operator< and is 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. check_for_forward_support). Added for convinence - // purposes and do not represent any graph properties. + // Internal fields: not keyed. The following fields are not part of attribute serialization, + // operator<, or the cache key; they are filled by derive() or set by caller such as with + // check_for_forward_support, and are used for convinence purposes // - // The two below say which directions nvte_get_fused_attn_backend_v2 should probe, and nothing - // else: not which graph is being built, which is what Pass names. They default to true and the - // attribute API cannot reach them, so a backend query from a framework asks about both - // directions, while the execution entry points set the one they are about to run. + // run query_support() for forward or backward bool check_for_forward_support = true; bool check_for_backward_support = true; - // Whether derive() has run, i.e. whether the fields below hold anything. Every consumer of a - // derived field needs them filled -- an unfilled config yields a graph with the wrong shapes - // and a cache key that collides with unrelated configs, neither of which announces itself -- - // so this exists to let those consumers assert rather than trust. Not a cached-result marker: - // derive() recomputes unconditionally, so a config whose inputs change can simply be re-derived. + // whether derive() has been run bool is_derived = false; - // THD batch/token counts, the raw buckets. The graph dimensions built out of them are - // graph_max_seqlen_* below and, because the batch is direction-dependent, graph_batch_size_*. + // bucketed batch size/token counts for THD 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. One answer per backend, because the same question has two: - // the FP8 graphs need newer cuDNN and frontend versions for it than the F16 ones, so a config - // that can hand cu_seqlens straight to one cannot necessarily hand them to the other. Each - // backend reads its own and no more; nothing reads both. + // whether to use cu_seqlens or actual_seqlens for THD or padding masks bool uses_cu_seqlens_directly = false; bool fp8_uses_cu_seqlens_directly = false; - // Whether a ragged (THD) graph is built at packed token-count dimensions with ragged Stats/LSE, - // rather than at dense max_seqlen ones. Held here rather than asked for at each of the places - // that need it -- graph_max_seqlen_* and graph_batch_size_* below -- because the key and the - // graph have to be built at the same dimensions, and two independent queries are two chances to - // disagree. Unlike the flags above, this one depends on the device as - // well as the cuDNN version, so a config carries the answer for the device it was derived on; - // every entry point derives immediately before use, and the cache key records device_id. + // Whether packed graphs exist for THD. A memory optimization, not a correctness gate: where this + // is false, ragged input still builds a correct graph, just the dense one, whose dimensions are + // max_seqlen rather than the token total and whose Stats is BHS1 rather than TH1. Nothing may + // gate support on it -- which architectures run ragged attention at all is cuDNN's answer. bool uses_packed_ragged_graph = false; - // Whether the graph's Stats/LSE tensor is the packed, token-indexed one. Ragged Q is necessary - // but not sufficient, since the packed representation also needs an architecture that supports - // it. Derived because three unrelated places read it -- the graph build, the pointer binding at - // execution, and the Stats/Max shapes reported back to the framework -- and they are describing - // one buffer, so they cannot be allowed to disagree about its shape. bool uses_ragged_stats = false; - // The sequence lengths the graph is built at: max_seqlen_* for a dense graph, and the bucketed - // token counts where a ragged layout is packed. Held here because the cache key has to name the - // dimensions the graph was built with -- a key that says otherwise is a hit on a graph of the - // wrong shape -- and stating the substitution once is what keeps make_cache_key() and the graph - // builders from drifting. Both passes build at the same sequence lengths; the batch size is the - // one dimension they disagree on, which is why that one is a pair below rather than a single - // field here. + // sequence lengths the graph is built at size_t graph_max_seqlen_q = 0; size_t graph_max_seqlen_kv = 0; - // The batch size the graph is built at and the width it expects ragged (THD) offsets in, one of - // each per direction. Pairs rather than one value apiece because the passes disagree on both, - // and a config is derived once and then probed and built for either direction, so no single - // field could answer: the selector derives a config and asks about forward and backward off that - // one copy. Whoever reads them names the direction they are building or keying for -- the graph - // builders, the code that binds runtime pointers to the built graph, and make_cache_key(), all - // of which have to agree, since a disagreement is a graph whose bound pointers do not describe - // the dimensions it was built at. All four are set together by the one condition in derive(). + // batch size the graph is built at size_t graph_batch_size_fwd = 0; size_t graph_batch_size_bwd = 0; + // ragged offset type for THD DType ragged_offset_type_fwd = DType::kInt32; DType ragged_offset_type_bwd = DType::kInt32; - // Elements per token for each ragged tensor, from the layout group and the head dimensions. - // Shared with the cu_seqlens_padded_to_offsets kernel, so the offsets the graph is told to - // expect and the offsets that are written cannot drift apart. + // Whether this config's ragged offsets overflow 32 bits. The counterpart to the two fields above + // rather than a third of them: those are the width TE will use, already capped by the running + // cuDNN, while this is the width the config needs. nvte_get_fused_attn_backend_v2 compares the + // two and refuses the config whose need outruns the cuDNN it is running on. + bool needs_64bit_ragged_offset = false; + // elements per token for each ragged tensor RaggedOffsetMultipliers ragged_offset_mults; - // Convinence fields to avoid recompute. + // convenience fields + // + // qkv_format is the combined format: it says what Q and KV each are and whether they agree, with + // the mixed layouts keeping their own enumerators (NVTE_THD_2BSHD and friends) rather than + // collapsing onto either side. Ask it only where a rule means "Q and KV are the same dense + // layout". Anything about raggedness belongs to is_ragged_q/is_ragged_kv instead, because + // NVTE_THD names only the fully ragged layouts, so a test against it passes THD_BSHD_BSHD + // straight through -- a rule here did exactly that until it was found. + NVTE_QKV_Format qkv_format = NVTE_QKV_Format_NOT_SET; 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; @@ -177,24 +145,29 @@ struct FusedAttnConfig { bool is_alibi = false; bool is_softmax_offset = false; bool is_mxfp8 = false; - // Whether the graph has a dropout node. The is_training term is what makes this one worth having - // as a field: a backward graph is only ever built for training, so the two directions used to - // spell this differently -- forward with the term, backward without -- and agreed only because - // every config that reaches a backward build has is_training set. One field states the rule the - // forward way, which is the safe way round: if that ever stops holding, a backward graph loses - // its dropout node rather than gaining one the cache key does not name. bool is_dropout = false; - // Whether what each pass stores is itself quantized: O for a forward graph, dQKV for a backward - // one. Only the FP8 backend asks, and for it this is the whole of what separates the two - // tensor-scaling recipes -- FP8 out means the scale is known before the graph is built, F16 out - // means the graph has to compute it. Derived rather than asked at each build site so that the - // pairing of a pass with the tensor it writes is stated once. + bool is_o_in_fp8 = false; + bool is_dqkv_in_fp8 = false; + // Whether the FP8 recipe is tensor scaling, i.e. delayed or current rather than MXFP8. The + // graphs need this on its own, wherever a tensor is per-tensor scaled and it does not matter + // which of the two put the scale there. + bool is_tensor_scaling = false; + // Which recipe serves each pass, one flag per recipe per pass. Each means "this recipe is in + // effect and can write this pass's output dtype", so at most one of a pass's three holds, and all + // three false is a configuration no FP8 graph is written for -- which is what lets + // nvte_get_fused_attn_backend_v2 refuse it by asking three booleans and nothing else. // - // The FP8 builders read "not FP8" as "F16", which holds only because - // nvte_get_fused_attn_backend_v2 refuses an FP8 config whose output is neither before any graph - // is built; that refusal and these two fields are the same rule read from its two ends. - bool o_is_fp8 = false; - bool dqkv_is_fp8 = false; + // Delayed against current is told apart by that output dtype rather than by scaling_mode, because + // NVTEScalingMode has no current-scaling enumerator: both arrive as NVTE_DELAYED_TENSOR_SCALING, + // and what separates them is that delayed knows the output scale before the graph is built while + // current has the graph compute it. So delayed writes FP8 and current writes F16/BF16, and an + // output dtype neither can write (FP32, say) leaves both false rather than defaulting to one. + bool is_delayed_scaling_fwd = false; + bool is_current_scaling_fwd = false; + bool is_mxfp8_fwd = false; + bool is_delayed_scaling_bwd = false; + bool is_current_scaling_bwd = false; + bool is_mxfp8_bwd = false; static constexpr size_t attr_sizes[] = { // basic attention settings @@ -247,11 +220,6 @@ struct FusedAttnConfig { sizeof(size_t), // bias_seqlen_kv }; - // The public header asks contributors to append to NVTEFusedAttnConfigAttribute, and the - // accessors index attr_sizes[attr] after checking only that attr is below the sentinel. An - // enumerator added without its size here would therefore read one past the end of this array, - // silently and only for the new attribute. Tying the two together turns that into a build - // failure at the line that has to change. static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnConfigNumAttributes, "attr_sizes must have one entry per NVTEFusedAttnConfigAttribute; add the size of " "the new attribute alongside its enumerator."); @@ -279,67 +247,33 @@ struct FusedAttnConfig { 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 - // that have been set by the caller. Call once, after the last input field is set and before - // the config reaches a graph build, a cache lookup, or a support query -- all of which read - // derived fields. - // - // Called by whoever owns the config, at the point it stops being edited: the execution entry - // points (nvte_fused_attn_fwd_v2 and its backward counterpart) on the config they go on to run, - // and nvte_get_fused_attn_backend_v2() on a copy of the caller's, so that asking whether a - // configuration is supported does not modify it. Nothing further in is expected to derive - // again, and check_derived() is what holds them to that. Idempotent, so a config that is - // derived and then re-derived is unharmed. + // Derive relevant fields based on input fields that have been set by the caller. They are + // read by the graph build, cache lookup, and support query. // - // Throws for combinations of input fields that no graph can serve, so that all four graph - // builders inherit the rule from one place. Those same combinations are stated as rejection - // rules in nvte_get_fused_attn_backend_v2(), ahead of its derive() call, so that asking whether - // such a configuration is supported gets an answer instead of an exception. + // Computes only; it validates nothing. Rules about which configurations are legal belong in + // nvte_get_fused_attn_backend_v2, which derives first and then states them once, so that a + // violation comes back as an unsupported configuration instead of being thrown from a query. void derive(); + // Assert that derive() has run, for code about to read a derived field. Worth asserting rather + // than assuming because the failure is silent: an unset derived field reads as zero, which is a + // legal value that yields a graph of the wrong shape and a key that collides with unrelated + // configs. + void check_derived() const { + NVTE_CHECK(is_derived, + "FusedAttnConfig's derived fields are not set. Please run " + "FusedAttnConfig::derive() first."); + } + // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. - // Requires a config that has been through derive(), whose fields the normalizations read. - // 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. - // - // `pass` is which graph the key is for. It decides both direction-dependent normalizations -- - // which fields are dropped, and whether the batch is bucketed -- so a key built for one pass - // cannot be handed to the other's cache. + // It drops fields that are either invariant (e.g. attn_scale) or irrelevant (e.g. dO/dQKV dtypes + // and `deterministic` for forward, and `return_max_logit` for backward). FusedAttnConfig make_cache_key(Pass pass) const; - // One line for the graph cache's level-2 trace: every field operator< compares, in its order, - // less device_id, which the trace's own prefix prints as the dev column. Abbreviated and terse - // on purpose, the reason to print a key at all being that two of these lines diff cleanly, - // naming the fields that cost an extra graph build. - // - // Four fields it prints that operator< does not compare, because a line without them cannot - // account for the values beside them: check_for_forward_support, which make_cache_key() sets - // from the pass and so is what says which direction's key this is, and the three bucketed_* - // inputs, which are what the normalization substituted into batch_size and the token counts. - // - // Defined here rather than with the diagnostics that print it because it is the third - // enumeration of these fields, after attr_sizes and operator< above. A field added to the key - // without being added here does not fail to build, it just stops appearing in the trace, so the - // three lists are kept where one change can see all of them. - std::string key_debug_string() const; + // Return a string representation of this config for level-2 cache diagnostics. + std::string to_string() const; }; -// Assert that `cfg` has been through derive(), for code about to read a derived field. Worth -// asserting rather than assuming because the failure is silent: an unset bucketed_batch_size or -// q_format reads as zero, which is a legal value that yields a graph of the wrong shape and a -// key that collides with unrelated configs. Deriving happens at the library's entry points rather -// than here, where it would be needed, so this is what keeps a new path into the builders from -// quietly skipping it. It catches a config that was never derived and nothing else: a config -// derived and then edited passes, so callers that change an input field re-derive rather than rely -// on this, which derive() being idempotent makes cheap. -inline void check_derived(const FusedAttnConfig &cfg) { - NVTE_CHECK(cfg.is_derived, - "FusedAttnConfig reached a graph build with its derived fields unset. Every config " - "must pass through FusedAttnConfig::derive() first; see the entry points in " - "fused_attn.cpp."); -} - inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); return reinterpret_cast(config); @@ -351,11 +285,19 @@ inline FusedAttnConfig *get_fused_attn_config_mutable(NVTEFusedAttnConfig config } struct FusedAttnFwdParams { + // Input tensors NVTETensor Q = nullptr; NVTETensor K = nullptr; NVTETensor V = nullptr; NVTETensor Bias = nullptr; NVTETensor SoftmaxOffset = nullptr; + // Intermediate tensors + NVTETensor S = nullptr; + // Output tensor + NVTETensor O = nullptr; + // Auxiliary context tensor pack + NVTETensorPack *Aux_CTX_Tensors = nullptr; + // Miscellaneous tensors NVTETensor cu_seqlens_q = nullptr; NVTETensor cu_seqlens_kv = nullptr; NVTETensor cu_seqlens_q_padded = nullptr; @@ -363,25 +305,24 @@ struct FusedAttnFwdParams { 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; + // Scalars + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; bool is_training = true; - bool cuda_graph = false; bool return_max_logit = false; - NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + bool cuda_graph = false; + float attn_scale = 1.0f; + float dropout = 0.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; 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; 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 qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; + // Workspace and stream NVTETensor workspace = nullptr; cudaStream_t stream = nullptr; @@ -391,6 +332,9 @@ struct FusedAttnFwdParams { sizeof(NVTETensor), // V sizeof(NVTETensor), // Bias sizeof(NVTETensor), // SoftmaxOffset + sizeof(NVTETensor), // S + sizeof(NVTETensor), // O + sizeof(NVTETensorPack *), // Aux_CTX_Tensors sizeof(NVTETensor), // cu_seqlens_q sizeof(NVTETensor), // cu_seqlens_kv sizeof(NVTETensor), // cu_seqlens_q_padded @@ -398,31 +342,26 @@ struct FusedAttnFwdParams { 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(uint8_t), // is_training - sizeof(uint8_t), // cuda_graph sizeof(uint8_t), // return_max_logit - sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(uint8_t), // cuda_graph + sizeof(float), // attn_scale + sizeof(float), // dropout + 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(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), // qkv_scale_inv_format - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv sizeof(NVTETensor), // workspace sizeof(cudaStream_t), // stream }; - // See FusedAttnConfig::attr_sizes: an enumerator appended without a size here reads past the - // end of this array. static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnFwdParamsNumAttributes, "attr_sizes must have one entry per NVTEFusedAttnFwdParamsAttribute; add the size " "of the new attribute alongside its enumerator."); @@ -445,6 +384,7 @@ inline FusedAttnFwdParams *get_fused_attn_fwd_params_mutable(NVTEFusedAttnFwdPar } struct FusedAttnBwdParams { + // Input tensors NVTETensor Q = nullptr; NVTETensor K = nullptr; NVTETensor V = nullptr; @@ -453,33 +393,37 @@ struct FusedAttnBwdParams { NVTETensor S = nullptr; NVTETensor dP = nullptr; const NVTETensorPack *Aux_CTX_Tensors = nullptr; + // Output tensors NVTETensor dQ = nullptr; NVTETensor dK = nullptr; NVTETensor dV = nullptr; NVTETensor dBias = nullptr; NVTETensor dSoftmaxOffset = nullptr; + // Miscellaneous tensors NVTETensor cu_seqlens_q = nullptr; NVTETensor cu_seqlens_kv = nullptr; NVTETensor cu_seqlens_q_padded = nullptr; NVTETensor cu_seqlens_kv_padded = nullptr; - 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; + // Scalars + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; float attn_scale = 1.0f; + float dropout = 0.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; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; + 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; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + bool deterministic = false; + bool cuda_graph = false; + // Workspace and stream NVTETensor workspace = nullptr; cudaStream_t stream = nullptr; @@ -501,30 +445,28 @@ struct FusedAttnBwdParams { sizeof(NVTETensor), // cu_seqlens_kv sizeof(NVTETensor), // cu_seqlens_q_padded sizeof(NVTETensor), // cu_seqlens_kv_padded - 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(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv sizeof(float), // attn_scale + sizeof(float), // dropout 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(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Mask_Type), // attn_mask_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(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph sizeof(NVTETensor), // workspace sizeof(cudaStream_t), // stream }; - // See FusedAttnConfig::attr_sizes: an enumerator appended without a size here reads past the - // end of this array. static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnBwdParamsNumAttributes, "attr_sizes must have one entry per NVTEFusedAttnBwdParamsAttribute; add the size " "of the new attribute alongside its enumerator."); diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index a3aea68a5c..dc8a5be41d 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -6,6 +6,8 @@ #include "transformer_engine/fused_attn.h" +#include + #include "../common.h" #include "../cudnn_utils.h" #include "../util/cuda_runtime.h" @@ -228,112 +230,65 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { namespace { -// The per-thread storage for the diagnostic string; it is re-used (cleared + re-populated) -// on every call to nvte_get_fused_attn_backend_v2 on the same thread. +// The per-thread storage for the diagnostic string, re-populated by every rejection on this +// thread. `*message` is handed a pointer into it, which is what limits how long that pointer +// stays good: only until the next rejection 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, -// 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); - *message = fused_attn_backend_message_buffer.c_str(); -} - // Records `reason` and answers with the backend that means "none", so that a rejection reads as -// the one statement it is: `if (cond) return reject(message, "why");`. Every rejection in -// nvte_get_fused_attn_backend_v2 goes through here, which is what keeps a reason attached to -// each: the value cannot be produced without one. nodiscard because dropping the value would -// leave the message set and the rejection unreturned, and the function would carry on. +// the one statement it is: `if (cond) return reject(message, "why");`. [[nodiscard]] NVTE_Fused_Attn_Backend reject(const char **message, std::string reason) { - set_message(message, std::move(reason)); + if (message != nullptr) { + fused_attn_backend_message_buffer = std::move(reason); + *message = fused_attn_backend_message_buffer.c_str(); + } return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } // namespace -// select a backend for fused attention; the diagnostic message is based on the first failure, not cumulative. +// Fused attention backend query: returns the backend that supports the given configuration; +// otherwise, returns NVTE_No_Backend and a diagnostic message. It performs several TE-specific +// checks before running the cuDNN support query. NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, const char **message) { NVTE_API_CALL(nvte_get_fused_attn_backend_v2); using namespace transformer_engine; using namespace transformer_engine::fused_attn; - // Derived on a copy, leaving the caller's config untouched: this function answers a question - // about a configuration and has no business editing one, and a query that wrote to its argument - // could not be asked about the same config from two threads at once. The copy costs nothing that - // matters here, since deriving is a version check and some arithmetic. - // - // The execution path derives its own config before calling this (see nvte_fused_attn_fwd_v2), - // and still reuses whatever graph the query builds: both derive the same fields from the same - // inputs, so make_cache_key() lands on the same entry. Deriving is idempotent, so re-deriving - // an already-derived config here changes nothing. - FusedAttnConfig cfg = *get_fused_attn_config(config); - set_message(message, ""); + const FusedAttnConfig &caller_cfg = *get_fused_attn_config(config); + if (message != nullptr) *message = ""; + std::optional derived_cfg; + if (!caller_cfg.is_derived) { + derived_cfg = caller_cfg; + derived_cfg->derive(); + } + const FusedAttnConfig &cfg = caller_cfg.is_derived ? caller_cfg : *derived_cfg; cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const auto qkv_format = nvte_get_qkv_format(cfg.qkv_layout); - const auto layout_group = nvte_get_qkv_layout_group(cfg.qkv_layout); const auto cudnn_runtime_version = cudnnGetVersion(); - // Read from attn_mask_type rather than from cfg.is_padding, because the two rules that need it - // are stated before derive() runs; see the derive() call below. - const bool has_padding_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; + const int sm_arch = cuda::sm_arch(cuda::current_device()); // 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); - if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { - return reject(message, - "Configuration requires 64-bit ragged offsets, which require cuDNN >= 9.5."); + if (cfg.needs_64bit_ragged_offset && cudnn_runtime_version < 90500) { + return reject( + message, + "This config requires 64-bit ragged offsets, which is only supported by cuDNN >= 9.5."); } - // THD requires padding-style mask - if (qkv_format == NVTE_QKV_Format::NVTE_THD && !has_padding_mask) { + // Ragged (THD) input requires a padding-style mask + if ((cfg.is_ragged_q || cfg.is_ragged_kv) && !cfg.is_padding) { return reject( message, "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); } - // Paged KV requires padding-style mask, for the same reason THD does: the graph is built at - // padded dimensions and the mask is what tells cuDNN where the real tokens end. - if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && !has_padding_mask) { + // Paged KV requires a padding-style mask, for the same reason ragged input does. + if (cfg.is_paged_kv && !cfg.is_padding) { return reject(message, "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); } - // Derived here rather than above, so that the two rules stated above are answered rather than - // thrown. Both are invariants derive() asserts, and an assertion that fired first would leave - // this function no chance to report them as an unsupported configuration. - cfg.derive(); - - // Ragged Q/KV requires sm90+, the rule the hand-written support matrix this function replaced - // carried as `qkv_format == NVTE_THD && sm_arch_ >= 90`. Below sm90 the only graph we can build - // is the dense max_seqlen one -- cfg.uses_packed_ragged_graph is false -- so SDPA_backward - // never gets max_total_seq_len_q/kv and its dQ/dK/dV come back wrong. - // - // This is ours to state because it is a wrong-result rejection, and check_support answers a - // different question: whether cuDNN can run the graph, not whether the graph computes what we - // asked for. cuDNN's own answer has moved, which is what makes the distinction worth spelling - // out here. Its frontend gates ragged SDPA on `sm < 90 && cudnn < 9.18.1`, so through 9.18.0 it - // would have refused this configuration for us and the rule below is redundant; from 9.18.1 it - // accepts sm80/sm89 ragged and the rule is the only thing standing between a THD model on an - // A100 and silently wrong gradients. Lifting it is a change to the graphs TE builds -- packed - // ragged shapes, and the Stats/LSE layouts that go with them, which cuDNN documents as - // differing on sm8x -- not a change to this condition, and that work is deliberately not part - // of this refactor. - // - // sm120 takes that same dense path and is left enabled, as it was before this refactor; - // whether it has the same problem is a separate question from restoring the sm90 rule. - if ((cfg.is_ragged_q || cfg.is_ragged_kv) && cuda::sm_arch(cuda::current_device()) < 90) { - return reject(message, "Ragged (THD) Q or KV requires compute capability 9.0 or higher."); - } - - // TE's cuDNN fused-attention graph does not represent pre-scale bias. + // cuDNN attention graphs do not support pre-scale bias if (cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) { return reject(message, "Fused attention does not support pre-scale bias."); } @@ -343,14 +298,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi const bool is_f16_or_bf16 = (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); - // Ask `verdict` about each direction the caller wants, and report the first refusal: the empty - // string means every direction asked about is served. Stated once here because every rule that - // is direction-dependent has to be asked this same way, and two copies of the gating would be - // two chances to probe a direction the caller never asked about. - // - // Forward is asked first because a config that cannot run forward cannot train either, and the - // forward refusal is the more useful of the two to report. Backward is skipped for inference, - // where no backward graph is ever built. + // Ask `verdict` about each direction the caller wants, and report the first refusal auto each_pass = [&](auto &&verdict) -> std::string { if (cfg.check_for_forward_support) { std::string reason = verdict(Pass::Fwd); @@ -363,55 +311,28 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi return ""; }; - // cuDNN's own verdict on `backend`. The two backends differ only in which set of graphs gets - // built, so they share this; what is theirs alone are the rules in each branch below. - // - // Each backend answers for both directions from its own translation unit, the only place that can - // name the graph builders, so choosing between them is all the dispatch left to do here. - auto probe = [&](Backend backend) -> std::string { - return each_pass([&](Pass pass) { - return backend == Backend::FP8 ? support_verdict_fp8(cfg, pass, handle) - : support_verdict_f16(cfg, pass, handle); - }); - }; - if (is_fp8) { if (cfg.return_max_logit) { return reject(message, "FP8 fused attention does not support return_max_logit=True."); } - if (qkv_format != NVTE_QKV_Format::NVTE_BSHD && qkv_format != NVTE_QKV_Format::NVTE_SBHD && - qkv_format != NVTE_QKV_Format::NVTE_BHSD) { + if (cfg.qkv_format != NVTE_QKV_Format::NVTE_BSHD && + cfg.qkv_format != NVTE_QKV_Format::NVTE_SBHD && + cfg.qkv_format != NVTE_QKV_Format::NVTE_BHSD) { return reject(message, "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + - std::to_string(static_cast(qkv_format)) + "."); + std::to_string(static_cast(cfg.qkv_format)) + "."); } - // The rest of what the FP8 graphs cannot represent: bias, ALiBi, and the quantization recipes - // they are not written for. TE's rules rather than cuDNN's, and stated here rather than in the - // build path for the reason all the rules above are: a rejection stated here is an answer - // carrying its reason, where the same rule inside a graph build would have to travel out as an - // exception. - if (cfg.bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { + if (cfg.is_bias) { return reject(message, "FP8 fused attention does not support pre/post_scale_bias yet!"); } - if (cfg.bias_type == NVTE_Bias_Type::NVTE_ALIBI) { + if (cfg.is_alibi) { return reject(message, "FP8 fused attention does not support ALiBi yet!"); } - // Whether the config names a recipe the FP8 graphs are written for at all. Delayed scaling - // writes FP8 out and keeps its scale, current scaling writes F16 and computes one, MXFP8 writes - // F16 with block scales; every other pairing of scaling mode and output dtype is refused here. - // - // Per direction, because the pairing is read off what each pass stores, which is also how the - // graph builders read which of the three they are building for -- off cfg.o_is_fp8 and - // cfg.dqkv_is_fp8, taking "not FP8" to mean F16. That reading is sound only because this - // refusal has already happened, so the two belong to each other: a pairing accepted here must - // be one they read the same way, and anything added to either belongs in both. std::string recipe_reason = each_pass([&](Pass pass) -> std::string { - const NVTEDType out_dtype = (pass == Pass::Fwd) ? cfg.o_dtype : cfg.dqkv_dtype; - const bool out_is_fp8 = (pass == Pass::Fwd) ? cfg.o_is_fp8 : cfg.dqkv_is_fp8; - const bool out_is_f16 = (out_dtype == kNVTEFloat16 || out_dtype == kNVTEBFloat16); const bool serves_this_output = - (cfg.scaling_mode == NVTE_DELAYED_TENSOR_SCALING && (out_is_fp8 || out_is_f16)) || - (cfg.scaling_mode == NVTE_MXFP8_1D_SCALING && out_is_f16); + (pass == Pass::Fwd) + ? (cfg.is_delayed_scaling_fwd || cfg.is_current_scaling_fwd || cfg.is_mxfp8_fwd) + : (cfg.is_delayed_scaling_bwd || cfg.is_current_scaling_bwd || cfg.is_mxfp8_bwd); if (!serves_this_output) { return "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 " "recipes!"; @@ -420,34 +341,37 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig confi }); if (!recipe_reason.empty()) return reject(message, std::move(recipe_reason)); - // Asked after the pairing above, not with it: a config that names no recipe at all should hear - // that rather than be sent to upgrade cuDNN for a recipe it was not asking for. - if (cfg.scaling_mode == NVTE_MXFP8_1D_SCALING && cudnn_runtime_version < 92100) { + if (cfg.is_mxfp8 && cudnn_runtime_version < 92100) { return reject(message, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); } - std::string reason = probe(Backend::FP8); - if (!reason.empty()) return reject(message, std::move(reason)); + std::string cudnn_reason = + each_pass([&](Pass pass) { return support_verdict_fp8(cfg, pass, handle); }); + if (!cudnn_reason.empty()) return reject(message, std::move(cudnn_reason)); return NVTE_Fused_Attn_Backend::NVTE_FP8; } if (is_f16_or_bf16) { - // TODO(cyanguwa): re-validate BRCM + cross-attention on sm100 with cuDNN <= 9.7. The - // hand-written support matrix this function replaced rejected bottom-right-diagonal masks - // with max_seqlen_q != max_seqlen_kv there, for a cuDNN bug fixed in 9.7. cuDNN's own - // check_support is the authority now, so the guard is gone; it needs to come back as an - // explicit rejection here, like the CUDA-graph one below, if that bug is a wrong-result - // bug rather than a support gap check_support reports for itself. + const bool has_sliding_window = !(cfg.window_size_left == -1 && + (cfg.window_size_right == -1 || cfg.window_size_right == 0)); + if (cfg.is_causal_bottom_right && has_sliding_window && cfg.max_seqlen_q != cfg.max_seqlen_kv && + cudnn_runtime_version <= 90700 && sm_arch >= 100) { + return reject(message, + "Known cuDNN <= 9.7.0 issue with bottom-right causal masking and a sliding " + "window for cross-attention on sm100. Please upgrade cuDNN."); + } + 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.qkv_format == NVTE_QKV_Format::NVTE_BSHD || + cfg.qkv_format == NVTE_QKV_Format::NVTE_SBHD) && + (cfg.max_seqlen_kv % 128 != 0) && cfg.cuda_graph && !cfg.is_padding) { return reject(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); } - std::string reason = probe(Backend::F16); - if (!reason.empty()) return reject(message, std::move(reason)); + + // run cudnn support checks + std::string cudnn_reason = + each_pass([&](Pass pass) { return support_verdict_f16(cfg, pass, handle); }); + if (!cudnn_reason.empty()) return reject(message, std::move(cudnn_reason)); return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } @@ -461,6 +385,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) { + NVTE_API_CALL(nvte_get_fused_attn_backend); transformer_engine::fused_attn::FusedAttnConfig cfg{}; cfg.qkv_layout = qkv_layout; cfg.bias_type = bias_type; @@ -484,7 +409,8 @@ 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 + // fill in the missing fields with the most common use case; + // otherwise it would return NVTE_No_Backend always cfg.batch_size = 1; cfg.o_format = nvte_get_q_format(qkv_layout); cfg.do_format = cfg.o_format; @@ -503,36 +429,20 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // Fused attention forward: derive the config, ask the selector which backend can run it, and run // that backend's implementation. // -// Support is decided by building the graph rather than by consulting a table of rules, and the -// support query and the execution path reach the same cache through the same accessor. That is -// what the HIT below means: by the time a backend has been selected, the entry the implementation -// needs has already been built and inserted by the probe that selected it, so what was checked is -// what runs. The rules the selector does state for itself are the ones cuDNN cannot answer: either -// about whether the graph computes what was asked for rather than whether cuDNN can run it, or -// about what TE's graphs can represent in the first place, which is where the FP8 recipes come in. -// Stating them here rather than inside a build is what lets each one answer with its reason. +// Both reach the same cache through the same accessor, which is what the HIT below means -- the +// entry execution needs was built by the probe that selected the backend, so what was checked is +// what runs. See graph_cache.h for the cache, graph_cache_debug.h for the events it emits. // // nvte_fused_attn_fwd_v2 // | -// +-- cfg = p.make_config(), which sets check_for_forward_support; cfg.derive() +// +-- cfg = p.make_config(); cfg.derive() // | -// +-- nvte_get_fused_attn_backend_v2 the support query -// | | -// | +-- TE's own rules: THD and paged KV need a padding mask, no pre-scale bias, -// | | ragged Q/KV needs sm90+, the cuDNN 9.15-and-older CUDA-graph bug, and for FP8 -// | | bias, ALiBi and the quantization recipes its graphs are not written for -// | | `-- reject -> NVTE_No_Backend + reason -> the NVTE_ERROR below -// | | -// | `-- probe(backend) -> support_verdict_f16 / support_verdict_fp8, with Pass::Fwd -// | `-- support_verdict() -// | `-- get_graph(): builds and inserts the entry, or throws, -// | in which case cuDNN's message becomes the reason for the refusal +// +-- nvte_get_fused_attn_backend_v2 TE's own rules, then a probe per backend +// | `-- support_verdict_f16 / support_verdict_fp8, with Pass::Fwd +// | `-- get_graph(): builds and inserts the entry, or refuses with cuDNN's reason // | -// `-- fused_attn_arbitrary_seqlen_fwd -> ..._fwd_impl the selected backend -// | -// +-- get_graph() HIT: the entry the query above just built -// +-- build_plans() the kernel compilation, once per entry -// `-- bind device pointers, graph.execute() +// `-- fused_attn_arbitrary_seqlen_fwd -> ..._fwd_impl the selected backend +// `-- get_graph() HIT -> build_plans() -> bind device pointers, graph.execute() void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { NVTE_API_CALL(nvte_fused_attn_fwd_v2); using namespace transformer_engine; @@ -556,8 +466,6 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); FusedAttnConfig cfg = p.make_config(); - // Derived here, not by the query below: the query works on its own copy, and it is this config - // that goes on to the backend and must arrive with its derived fields filled in. cfg.derive(); const char *fused_attn_reject_reason = nullptr; NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( @@ -597,7 +505,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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); + NVTE_API_CALL(nvte_fused_attn_fwd); transformer_engine::fused_attn::FusedAttnFwdParams p{}; p.Q = Q; p.K = K; @@ -635,9 +543,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso nvte_fused_attn_fwd_v2(reinterpret_cast(&p)); } -// Fused attention backward. The same shape as nvte_fused_attn_fwd_v2, which sketches the path from -// an entry point through the selector to the cache; the only differences are that the config asks -// the selector for backward support and that the backward builders are the ones the probe runs. +// Fused attention backward. Same shape as nvte_fused_attn_fwd_v2, whose comment sketches the path; +// this one asks the selector for backward support and probes the backward builders. void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { NVTE_API_CALL(nvte_fused_attn_bwd_v2); using namespace transformer_engine; @@ -727,7 +634,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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_API_CALL(nvte_fused_attn_bwd); transformer_engine::fused_attn::FusedAttnBwdParams p{}; p.Q = Q; p.K = K; 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 e68595fd26..33f0ef52a1 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 @@ -375,7 +375,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // and the ones the graph was built at cannot be decided differently. Asserted derived here // because these are the first derived fields this path reads, ahead of the get_graph() that // asserts it for the build. - check_derived(cfg); + cfg.check_derived(); const int64_t b = static_cast(cfg.graph_batch_size_fwd); const DType ragged_offset_type = cfg.ragged_offset_type_fwd; // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever the @@ -853,7 +853,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( // and the ones the graph was built at cannot be decided differently. Asserted derived here // because these are the first derived fields this path reads, ahead of the get_graph() that // asserts it for the build. - check_derived(cfg); + cfg.check_derived(); const int64_t b = static_cast(cfg.graph_batch_size_bwd); const DType ragged_offset_type = cfg.ragged_offset_type_bwd; // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by. diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 91361f1d0a..eeec53207e 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -44,18 +44,30 @@ using Fp8FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The three recipes these graphs are written for, spelled the same way at each of the four sites -// that build or bind one: +// The three recipes these graphs are written for, read from cfg at each of the four sites that +// build or bind one: // -// is_mxfp8 = scaling_mode is MXFP8 -// is_delayed_scaling = !is_mxfp8 && -// is_current_scaling = !is_mxfp8 && ! +// is_mxfp8 = scaling_mode is MXFP8 +// is_tensor_scaling = scaling_mode is DELAYED_TENSOR_SCALING +// is_delayed_scaling_fwd / _bwd = is_tensor_scaling && O / dQKV is FP8 +// is_current_scaling_fwd / _bwd = is_tensor_scaling && O / dQKV is F16 +// is_mxfp8_fwd / _bwd = is_mxfp8 && O / dQKV is F16 // -// which is exactly one of the three by construction, no combination of the booleans being able to -// say two things at once. The output half comes from cfg.o_is_fp8 or cfg.dqkv_is_fp8 -- a forward -// graph writes O, a backward one dQKV -- and reads "not FP8" as F16, which holds because -// nvte_get_fused_attn_backend_v2 refuses an FP8 config whose output is neither. See there for the -// rest of what these graphs cannot represent, and config_and_params.h for the two fields. +// so at most one of the three holds for a pass, no combination of the booleans being able to say +// two things at once, and is_tensor_scaling is the delayed/current pair together -- which is what +// most of the sites below want, since a per-tensor scale is a per-tensor scale whichever recipe +// put it there. +// +// Each flag pairs a recipe with an output dtype it can write, so an output none of them can write +// leaves all three false rather than defaulting to one. That is the form the check takes in +// nvte_get_fused_attn_backend_v2, which refuses such a config before any graph here is built -- +// which is in turn why the sites below can treat the three as a partition. +// +// The split is per pass because a forward graph writes O and a backward one dQKV, and it is drawn +// on the output dtype because NVTEScalingMode has no current-scaling enumerator: both +// tensor-scaling recipes arrive as DELAYED_TENSOR_SCALING, and what separates them is whether the +// scale is known before the graph is built. See nvte_get_fused_attn_backend_v2 for the rest of +// what these graphs cannot represent, and config_and_params.h for the fields. // // Unlike the F16 path there is no bucketing to do, because FP8 has no ragged/THD support: the // graph's shapes are exactly the config's. @@ -96,8 +108,9 @@ static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { const bool is_dropout = cfg.is_dropout; const bool is_softmax_offset = cfg.is_softmax_offset; const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_delayed_scaling = !is_mxfp8 && cfg.o_is_fp8; - const bool is_current_scaling = !is_mxfp8 && !cfg.o_is_fp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; + const bool is_current_scaling = cfg.is_current_scaling_fwd; const bool use_cu_seqlens_directly = cfg.fp8_uses_cu_seqlens_directly; auto mha_graph = std::make_shared(); @@ -138,7 +151,7 @@ static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { .set_data_type(fe::DataType_t::FLOAT)); // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Descale_q") .set_dim({1, 1, 1, 1}) @@ -280,7 +293,7 @@ static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { } std::shared_ptr O, Stats, amax_s, amax_o; - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, sdpa_options); O = outputs[0]; @@ -354,13 +367,12 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de // Asserted derived here because the reads below are the first derived fields this path touches, // ahead of the get_graph() that asserts it for the build. - check_derived(cfg); + cfg.check_derived(); // Read from the same fields the graph was built from, so that the tensors bound below and the // ones the graph was built with cannot be decided differently. - const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_delayed_scaling = !is_mxfp8 && cfg.o_is_fp8; - const bool is_current_scaling = !is_mxfp8 && !cfg.o_is_fp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; const bool use_cu_seqlens_directly = cfg.fp8_uses_cu_seqlens_directly; const int64_t b = static_cast(cfg.batch_size); @@ -408,7 +420,7 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de if (is_delayed_scaling) { variant_pack[scale_o] = devPtrScaleO; } - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { variant_pack[descale_s] = devPtrDescaleS; variant_pack[scale_s] = devPtrScaleS; variant_pack[amax_s] = devPtrAmaxS; @@ -543,11 +555,12 @@ static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { const bool is_dropout = cfg.is_dropout; const bool is_softmax_offset = cfg.is_softmax_offset; const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_delayed_scaling = !is_mxfp8 && cfg.dqkv_is_fp8; - const bool is_current_scaling = !is_mxfp8 && !cfg.dqkv_is_fp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; + const bool is_current_scaling = cfg.is_current_scaling_bwd; // Whether O arrived in F16 rather than FP8, which decides whether this graph has to descale it on // the way in. Read off O, unlike the recipe above, because O is what the forward pass stored. - const bool is_O_in_F16 = !cfg.o_is_fp8; + const bool is_O_in_F16 = !cfg.is_o_in_fp8; auto mha_graph = std::make_shared(); @@ -611,7 +624,7 @@ static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { .set_data_type(fe::DataType_t::FLOAT)); // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Descale_q") .set_dim({1, 1, 1, 1}) @@ -829,7 +842,7 @@ static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { } std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, @@ -870,7 +883,7 @@ static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { amax_dP->set_output(true) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) @@ -938,14 +951,15 @@ void fused_attn_fp8_bwd_impl( // Asserted derived here because the reads below are the first derived fields this path touches, // ahead of the get_graph() that asserts it for the build. - check_derived(cfg); + cfg.check_derived(); // Read from the same fields the graph was built from, so that the tensors bound below and the // ones the graph was built with cannot be decided differently. const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_delayed_scaling = !is_mxfp8 && cfg.dqkv_is_fp8; - const bool is_current_scaling = !is_mxfp8 && !cfg.dqkv_is_fp8; - const bool is_O_in_F16 = !cfg.o_is_fp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; + const bool is_current_scaling = cfg.is_current_scaling_bwd; + const bool is_O_in_F16 = !cfg.is_o_in_fp8; const int64_t b = static_cast(cfg.batch_size); // Not const: bound into the variant pack by address as a pass-by-value graph input. @@ -994,7 +1008,7 @@ void fused_attn_fp8_bwd_impl( {dK, devPtrdK}, {dV, devPtrdV}, }; - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { variant_pack[descale_s] = devPtrDescaleS; variant_pack[descale_dP] = devPtrDescaledP; variant_pack[scale_s] = devPtrScaleS; diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index 521dcf6748..21c5098ab5 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -4,31 +4,28 @@ * See LICENSE for license information. ************************************************************************/ -// ============================================================================ -// The fused-attention graph cache: what a cache entry is, how one is looked up -// or built, and the frontend calls that make a constructed graph usable. +// The fused-attention graph cache: what a cache entry is, and how one is built, cached and found +// again. The four build sites for cuDNN graphs -- f16/fp8 crossed with fwd/bwd -- differ only in +// what the graph computes and which tensors it binds; caching, lookup, locking, the support query +// and the plan build are the same for all four and live here. // -// The four build sites (f16 and fp8, forward and backward) differ only in how -// they construct their graph and which tensors they hand back. Everything after -// that -- the cache each one keeps, the lookup, the locking, the support check, -// the once-per-entry plan build -- is shared, and lives here rather than in four -// copies: a site names its backend, pass and graph builder to get_graph(). +// The pieces below elide the `backend, pass` pair most of them also take: it never steers the +// logic, only attributing debug counters and stage timings to a build site. // -// The five frontend calls a graph goes through, and which of our functions pays for -// each. The frontend's are written graph.*, since that is how they are invoked and -// since two of them share a name with ours: -// -// on a miss, either caller: -// graph.validate() -> graph.build_operation_graph() -// -> graph.create_execution_plans(HeurMode_t::A) -> graph.check_support() -// all four via query_support() -// the execution path only: -// graph.build_plans() build_plans(), once per entry, the kernel compilation -// graph.execute() every call, with its variant pack built in a local -// -// Not part of utils.h: this needs the cuDNN frontend, and utils.h is included by -// translation units (utils.cu) that otherwise do not. -// ============================================================================ +// - CacheEntry: a graph, the tensors it binds as inputs and outputs, and a once_flag guarding its +// plan build. +// - GraphCache: process-wide map from a normalized FusedAttnConfig to a CacheEntry. +// - get_graph(cfg, handle): the execution path's way in. Keys `cfg` +// and owns the cache for its one triple; kCreateGraphFn is a create_graph_f16/fp8_fwd/bwd from a +// .cu file, the only piece a build site supplies. +// - support_verdict<...>(cfg, handle): the backend selector's way in. get_graph() in a try, +// returning the empty string when cuDNN accepts the graph and its complaint when it does not. +// - cache_graph(cache, key, handle, build): a hit, or a build under frontend_build_mutex() and an +// insert. The work behind both of the above. +// - query_support(graph, handle): takes a constructed graph through validate, +// build_operation_graph, create_execution_plans and check_support; throws cuDNN's message on +// refusal. +// - build_plans(entry): the kernel compilation cache_graph() deferred, once per entry, no handle. #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ @@ -52,16 +49,11 @@ namespace fused_attn { // A graph in the cache, plus the tensor attributes needed to bind runtime pointers to it. // -// Entries are built only as far as check_support(), which is all it takes to decide whether a -// configuration is supported. graph.build_plans() -- the kernel compilation, and the most expensive -// of the five frontend calls -- is left to the execution path, since a support query never runs the -// graph and many of the keys it builds are never run by anything. -// -// build_plans_once guards that completion, which has to happen exactly once per entry: the entry is -// shared across threads and graph.build_plans() mutates it in place. Keeping the flag in the entry -// keeps it with the graph it describes and leaves unrelated keys free to build concurrently. A -// build that throws leaves it unset, so a later call retries rather than executing a graph with no -// plans. +// Entries stop at check_support(), which is all it takes to decide support. graph.build_plans() -- +// the kernel compilation, and the most expensive frontend call -- is left to the execution path, +// since a support query never runs the graph. build_plans_once guards that completion, which must +// happen exactly once per entry: the entry is shared across threads and graph.build_plans() mutates +// it in place. A build that throws leaves the flag unset, so a later call retries. template struct CacheEntry { explicit CacheEntry(GraphAndTensors graph_and_tensors) @@ -71,47 +63,48 @@ struct CacheEntry { std::once_flag build_plans_once; }; -// One build site's cache. Process-wide rather than per-thread so a graph is reused across threads +// One build site's cache, process-wide rather than per-thread so a graph is reused across threads // instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a shared plan, and the -// frontend's execute() builds its variant pack in a local rather than in the graph, so it does not -// write to the shared object. +// frontend's execute() builds its variant pack in a local rather than in the graph. A graph and its +// plans are compiled artifacts bound to the device they were finalized against, with nothing in +// them belonging to the building thread, which is why the key stamps device_id and nothing +// thread-shaped (see make_cache_key). // -// What lets one cache serve every thread is an asymmetry between the two objects a call needs. A -// cuDNN handle is per-thread mutable session state: it carries the stream execute() launches on, so -// each thread holds its own. A graph and its plans are the opposite -- compiled artifacts, bound to -// the device they were finalized against, with nothing in them belonging to the building thread. So -// the key stamps device_id and nothing thread-shaped (see make_cache_key), and build_plans() below -// covers the seam where the thread that finishes a build is not the one that started it. +// The mutex is declared first so that it is destroyed last, leaving the map destroyed while its +// guard is still valid. // -// The mutex is declared first so that it is destroyed last -- members go in reverse declaration -// order, so the map is destroyed while its guard is still valid. Declaring the two together settles -// that rather than leaving it to whoever writes the next cache. -// -// The map is unbounded. Only executed graphs hold anything substantial -- an entry that stopped at -// check_support() has no compiled kernels behind it, and none hold a workspace, which the caller -// allocates per call -- and a model reuses a handful of configurations, so any bound worth setting -// would sit far above what real work reaches. A workload that does sweep shapes, such as a suite -// enumerating them, holds every graph for the life of the process; `miss` climbing without settling -// is what that looks like, and is the case for bringing a bound back. +// The map is unbounded: a probe-only entry holds no compiled kernels, none hold a workspace, and a +// model reuses a handful of configurations. A workload that does sweep shapes holds every graph for +// the life of the process, which `miss` climbing without settling is the way to see. template struct GraphCache { std::mutex mutex; // guards everything below std::map>> entries; }; -// Takes a constructed graph through the frontend calls that decide whether cuDNN can run it: -// validate, build_operation_graph, create_execution_plans, check_support. Identical for both passes -// and both backends, so it is defined once here; `backend` and `pass` only name the build site the -// stage timers attribute the calls to. +// Every cuDNN frontend call except graph.execute() runs holding this. The frontend serializes none +// of them for us, so two threads building unrelated keys is a data race, not the harmless duplicate +// work the map's view suggests. Not a theoretical exposure: a PyTorch step runs the forward and the +// backward's support probe on the main thread and the backward itself on the autograd thread. // -// Reports by throwing, and the throw carries cuDNN's message alone. That message is what -// support_verdict() returns as the reason a backend was refused, so a bool would discard the one -// thing a support probe exists to produce -- and NVTE_ERROR would wrap it in the file, line and -// advice of an internal failure, which a backend refused for a plain reason is not. One kind of -// throw for every failure; see support_verdict() for why that distinction is not drawn. +// One lock for the process rather than one per cache, since what is unsafe is the frontend rather +// than any single graph. Kept separate from GraphCache::mutex, which guards only the map, so a hit +// never waits behind somebody else's kernel compilation. // -// graph.build_plans() and graph.execute() sit outside this function: they commit real resources, -// and the plan build belongs to whoever executes the graph, once. See CacheEntry. +// Lock ordering, which a later edit has to preserve: always taken before GraphCache::mutex, never +// after, and never held on entry to build_plans() -- holding it while waiting on that once_flag +// would deadlock against the thread holding the flag. +inline std::mutex &frontend_build_mutex() { + static std::mutex mutex; + return mutex; +} + +// Takes a constructed graph through the frontend calls that decide whether cuDNN can run it. +// `backend` and `pass` only name the build site the stage timers attribute the calls to. +// +// Reports by throwing, carrying cuDNN's message alone: that message is what support_verdict() +// returns as the reason a backend was refused, so a bool would discard the one thing a probe exists +// to produce, and NVTE_ERROR would dress a plain refusal as an internal failure. inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { @@ -133,38 +126,23 @@ inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Gra [&] { return graph.check_support(); }); } -// The cached entry for `key`, building and inserting it via `build` if absent. Throws if cuDNN -// refuses the graph, and remembers nothing when it does, so the next query for a refused key builds -// it again and is refused again. The frameworks only re-enter the selector when the attention -// configuration changes (in PyTorch, _attention_backends caches the choice), so a settled run pays -// for a refusal once; a suite that enumerates configurations pays each time it comes back around, -// which is the case a map of remembered refusals would serve. +// The cached entry for `key`, building and inserting it via `build` if absent: // -// `build` only constructs a graph; this is what puts it through query_support(), so the entries in -// the cache are exactly the graphs cuDNN has agreed to run. Those calls belong to building an entry -// rather than reading one, which is why a hit skips them. +// hit -> record HIT, return the entry +// miss -> take frontend_build_mutex(), look again (a thread that raced us has finished by now), +// record MISS, build(), query_support(), insert // -// `key` must be make_cache_key(pass)'s output, for the same `pass` given here, and not a raw -// execution config: two configs differing only in a field no graph reads (attn_scale, say) have to -// reach the same entry, and passing the raw config silently multiplies the cache by fields the -// graph never consumes. +// `build` only constructs a graph; this is what puts it through query_support(), so the cache holds +// exactly the graphs cuDNN agreed to run, and a hit skips those calls. A refusal throws and stores +// nothing, so the next query for a refused key is refused again -- fine for a settled run, since +// the frameworks re-enter the selector only when the configuration changes. // -// Only the map operations are locked, not `build`, so builds of unrelated keys proceed concurrently -// and two threads racing on one key may both build. That is wasted work rather than a correctness -// problem -- the loser drops its graph and takes the winner's entry, whose once_flag still governs -// the plan build -- and it reads in diagnostics as two MISS lines with the same key. +// `key` must be make_cache_key(pass)'s output for the same `pass`, not a raw execution config: two +// configs differing only in a field no graph reads (attn_scale, say) have to reach the same entry. // -// lock cache.mutex -// entries[key]? found -> copy the shared_ptr -// unlock -// record_hit_miss(HIT | MISS) -// -// HIT -> return the entry -// MISS -> build(), record_create_graph() outside the lock, so builds of -// query_support() unrelated keys proceed concurrently -// ok -> record_cache_graph(), then lock, insert, unlock; return the inserted -// entry, which on a lost race is the winner's -// throw -> propagates; nothing is stored, so the key is built again if it comes back +// The second look keeps a lost race cheap -- the loser would otherwise hold the one build lock to +// produce a graph it drops on the next line -- and keeps exactly one HIT or MISS per call, so miss +// still equals create_graph. template std::shared_ptr> cache_graph(GraphCache &cache, const FusedAttnConfig &key, @@ -172,97 +150,80 @@ std::shared_ptr> cache_graph(GraphCache> cached; - { + auto find = [&]() -> std::shared_ptr> { std::lock_guard lock(cache.mutex); auto it = cache.entries.find(key); - if (it != cache.entries.end()) cached = it->second; - } + return it != cache.entries.end() ? it->second : nullptr; + }; + // Recorded after the lock is dropped, so writing a trace line cannot hold up threads querying // other keys. The counters stay exact, but two lookups that raced can be recorded in the opposite - // order, so a level-2 trace is the set of lookups that happened, not their sequence. - graph_cache_debug::record_hit_miss( - backend, pass, cached != nullptr ? LookupResult::Hit : LookupResult::Miss, key); - if (cached != nullptr) return cached; + // order, so a level-2 trace is the set of lookups, not their sequence. + if (std::shared_ptr> cached = find()) { + graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Hit, key); + return cached; + } + + std::lock_guard build_lock(frontend_build_mutex()); + if (std::shared_ptr> cached = find()) { + graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Hit, key); + return cached; + } + // The one trace line written under the build lock. A build dwarfs an fprintf, and recording the + // miss before the lock would count a raced key as both a miss and a hit. + graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Miss, key); - // No backend refuses a configuration from in here any more -- TE's own FP8 rules moved to - // nvte_get_fused_attn_backend_v2, which answers them with a reason instead of throwing -- so - // `build` returning is now the ordinary case and miss == create_graph in a run that behaves. - // A throw that does get out still propagates with its message and leaves nothing behind, which - // is what keeps miss - create_graph worth printing: it should read zero. auto entry = std::make_shared>(build()); graph_cache_debug::record_create_graph(backend, pass); - // Every site's tensor tuple leads with its graph, the one element this needs. A tuple ordered - // otherwise would fail to compile rather than quietly validate the wrong object. + // Every site's tensor tuple leads with its graph, the one element this needs; a tuple ordered + // otherwise fails to compile rather than quietly validating the wrong object. // // The two counters bracket this call deliberately: a graph cuDNN refuses throws here, having - // already recorded its CREATE_GRAPH and never reaching CACHE_GRAPH, so the gap between those two - // columns is cuDNN's refusals alone. + // recorded its CREATE_GRAPH and never reaching CACHE_GRAPH, so the gap between those two columns + // is cuDNN's refusals alone. query_support(backend, pass, *std::get<0>(entry->graph_and_tensors), handle); - // Recorded on cuDNN's verdict rather than on the insert below, so the column counts the graphs - // cuDNN agreed to run. That is the question worth a counter; how many entries a map ended up - // holding is not, and tying it to the insert made a lost race -- which discards a supported graph - // and takes the winner's -- read as a miscount rather than as the duplicate work it is. graph_cache_debug::record_cache_graph(backend, pass); + // The insert always takes: a thread racing this key would have had to hold the build lock to do + // it, and the look above already ruled that out. std::lock_guard lock(cache.mutex); - // On a losing race the insert does nothing: the shared_ptr this thread built is dropped with - // its graph, and what comes back is the winner's entry. return cache.entries.insert({key, std::move(entry)}).first->second; } -// A backend's graph cache for one pass, and the only route to it. Both the execution path and the -// support probe come through here, so a probe leaves behind exactly the entry a later execution -// finds. That is what lets the probe's answer describe the graph that actually runs, rather than a -// separately built lookalike. -// -// The cache is this instantiation's static local, so the callers that name one triple share one cache, and each triple gets its own. Naming the triple is now what -// picks the cache, where before there was a per-backend function per pass to call. +// A backend's graph cache for one pass, and the only route to it. The cache is this instantiation's +// static local, so the callers naming one triple share one cache and each +// triple gets its own. // -// `kCreateGraphFn` is a template parameter rather than a `CreateFn &&` argument on purpose. As a -// parameter it makes the creator part of the instantiation, keeping the cache identified by the -// function that fills it. Passed as an argument, each distinct lambda type would instantiate its -// own copy of this function with its own static cache, and the two call sites for a pass would +// `kCreateGraphFn` is a template parameter rather than a `CreateFn &&` argument on purpose: as a +// parameter it makes the creator part of the instantiation. Passed as an argument, each distinct +// lambda type would instantiate its own static cache, and the two call sites for a pass would // quietly stop sharing entries. template auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { static GraphCache cache; // Asserted once here for both the key and the graph, which read the same derived fields. - check_derived(cfg); + cfg.check_derived(); return cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, [&] { return kCreateGraphFn(cfg); }); } // Whether cuDNN can run the graph this config asks for, in one direction: the empty string if it -// can, otherwise cuDNN's own account of why not, which the backend selector reports to the caller. -// This is the whole of what support_verdict_f16 and support_verdict_fp8 do; they exist only to -// reach their own translation unit's graph builders, which is also where a runtime direction turns -// into the compile-time one this needs. -// -// Named for what it returns rather than the question it answers: support is the empty string, so -// an is_supported() spelling would read backwards wherever the result is tested. -// -// The question is answered by building the graph, which is where every rejection comes from -- -// there is no separate list of rules to keep in step with the builder. The graph goes into the same -// cache the execution path reads, so the work is not thrown away and what was checked is what will -// run. It stops short of graph.build_plans(), the expensive step, which the first execution of the -// graph does instead; see CacheEntry. +// can, otherwise cuDNN's own account of why not. This is the whole of what support_verdict_f16 and +// support_verdict_fp8 do; they exist only to reach their own translation unit's graph builders, +// which is also where a runtime direction becomes the compile-time one this needs. // -// A refusal, by contrast, is not cached: nothing is stored for a key cuDNN rejected, so asking the -// same question again pays for the build again. See cache_graph. +// The question is answered by building the graph, so there is no separate list of rules to keep in +// step with the builder, and the graph lands in the cache the execution path reads. // -// Refusals and failures on the way to a verdict read alike, because CUDNN_BACKEND_API_FAILED -- -// raised for any non-success cudnnStatus_t -- cannot separate CUDNN_STATUS_NOT_SUPPORTED from -// CUDNN_STATUS_ALLOC_FAILED. Either way this backend cannot serve this call, and either way what -// the caller wants is the message. +// Refusals and failures read alike, because CUDNN_BACKEND_API_FAILED -- raised for any non-success +// cudnnStatus_t -- cannot separate CUDNN_STATUS_NOT_SUPPORTED from CUDNN_STATUS_ALLOC_FAILED. +// Either way this backend cannot serve this call and the caller wants the message. // -// The direction is named by the caller rather than read off the config: a config arriving from a -// framework has both check_for_*_support set, so it cannot say which graph is being probed. +// The direction is named by the caller rather than read off the config, which has both +// check_for_*_support set and so cannot say which graph is being probed. template std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - // Built only where it is used, on the two paths where a refusal arrived without a message of its - // own. Support is signalled by returning the empty string, so an empty refusal would otherwise - // read as an endorsement. + // Support is signalled by returning the empty string, so a refusal that arrived without a message + // of its own needs a label rather than reading as an endorsement. auto label = [] { return std::string("support_verdict<") + graph_cache_debug::backend_name(kBackend) + ", " + graph_cache_debug::pass_name(kPass) + ">"; @@ -279,35 +240,29 @@ std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { } } -// Runs graph.build_plans(), the plan build that cache_graph() left undone, once per entry. Named -// for the frontend call it wraps; the once-per-entry part is the whole reason it is a function -// rather than that call. +// Runs graph.build_plans(), the plan build cache_graph() left undone, once per entry. Call only +// when the graph is about to be executed: a support query builds entries nothing ever runs, and +// kernel compilation is the most expensive frontend call. See CacheEntry. // -// Call only when the graph is about to be executed, which is why this is a separate step rather -// than the tail of the lookup: a support query builds entries nothing ever runs, and kernel -// compilation is the most expensive of the five frontend calls. See CacheEntry for why the flag -// lives inside the entry and what a throw here leaves behind. +// The once_flag settles which thread runs the build, not whether it may run alongside another one: +// graph.build_plans() is a frontend call like the rest, so it also needs frontend_build_mutex(). +// The lock is taken inside the call_once rather than around it, so an entry whose plans are already +// built stays on the flag's atomic fast path and never touches the process-wide lock. // -// Splitting the build in two means the thread that finishes it is often not the thread that started -// it -- a sizing call caches the graph, and an autograd thread is the first to need it to run. Four -// facts make that safe, and only the first is visible here: -// - graph.build_plans() takes no handle. The overload that accepts one ignores it (its body is -// `(void)handle;`), working from the operation graph descriptor and the device properties -// instead, which is how deviceless ahead-of-time compilation builds plans with no handle at -// all. Unlike the plan sharing on GraphCache, this does lean on the >= 1.25.0 frontend the -// build requires: it is where the handle-free overload arrived. -// - The handle that built the operation graph outlives the build, held by the descriptor -// graph.build_operation_graph(handle) finalized against it, and stays valid only because TE -// never destroys cuDNN handles: cudnnExecutionPlanManager leaves HandleManager's Destroy -// parameter at its nullptr default, so handles leak by design, one per thread per device. -// - That descriptor was finalized for one device, which is why the cache key carries device_id -// (see make_cache_key). Without it a thread could compile kernels from another device's -// descriptor. -// - graph.execute() is called with the running thread's own handle, so a handle is never used by -// two threads at once, which is what cuDNN asks in return for letting them share a plan. +// Splitting the build in two means the thread that finishes it is often not the one that started it +// -- a sizing call caches the graph, an autograd thread is first to run it. What makes that safe: +// graph.build_plans() needs no handle (the overload accepting one ignores it, working from the +// operation graph descriptor and the device properties, which is where the >= 1.25.0 frontend the +// build requires is load-bearing); the handle that built that descriptor outlives the build because +// TE never destroys cuDNN handles (cudnnExecutionPlanManager leaves HandleManager's Destroy +// parameter null, so handles leak by design); the descriptor was finalized for one device, which is +// why the key carries device_id; and graph.execute() uses the running thread's own handle, so a +// handle is never used by two threads at once, which is what cuDNN asks in return for letting them +// share a plan. template void build_plans(Backend backend, Pass pass, CacheEntry &entry) { std::call_once(entry.build_plans_once, [&] { + std::lock_guard build_lock(frontend_build_mutex()); cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.graph_and_tensors); graph_cache_debug::record_time(backend, pass, graph_cache_debug::BuildStage::BuildPlans, [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 882ec3b22a..8358e7e24d 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -4,26 +4,24 @@ * See LICENSE for license information. ************************************************************************/ -// ============================================================================ // Fused-attention graph cache diagnostics. // -// Enable with NVTE_FUSED_ATTN_CACHE_DEBUG=[:]. The output format, how -// to read it and the rank suffix are documented for users in docs/envvars.rst; what -// follows is what maintaining this file needs. +// Enable with NVTE_FUSED_ATTN_CACHE_DEBUG=[:]. The output format, how to read it and +// the rank suffix are documented for users in docs/envvars.rst; what follows is what maintaining +// this file needs. // -// level 1 (events) : one line per event that happens once per distinct cache key -// (CREATE_GRAPH, CACHE_GRAPH, BUILD_PLANS), plus the exit summary -// block and its stage timings. Low volume by construction. -// level 2 (trace) : adds a line per cache lookup (HIT/MISS, with the normalized -// key) and per execution (EXECUTE). High volume, and it serializes -// threads on the stderr lock, which the stage timings are then -// measured under -- no timed region writes to stderr, so they -// stay sound, but they read a little high. +// level 1 (events) : one line per event that happens once per distinct cache key (CREATE_GRAPH, +// CACHE_GRAPH, BUILD_PLANS), plus the exit summary and its stage timings. +// level 2 (trace) : adds a line per lookup (HIT/MISS, with the normalized key) and per execution +// (EXECUTE). High volume, and it serializes threads on the stderr lock, which +// the stage timings are then measured under -- no timed region writes to +// stderr, so they stay sound, but they read a little high. // -// Counters are kept per build site -- f16/fp8 crossed with fwd/bwd -- since one -// process can drive both backends, and every event name is also the counter column -// it increments. What the columns mean, the identities they satisfy and the ratios -// worth reading are with the counter definitions below. +// Counters are kept per build site -- f16/fp8 crossed with fwd/bwd -- since one process can drive +// both backends, and every event name is also the counter column it increments. Where the events +// sit on the path nvte_fused_attn_fwd_v2 sketches: HIT/MISS on every get_graph() lookup, +// CREATE_GRAPH and CACHE_GRAPH inside it on a miss, BUILD_PLANS on an entry's first execution, and +// EXECUTE on every call. // // One level-1 training step, line prefixes and trailing columns elided: // @@ -36,22 +34,12 @@ // f16 fwd build_plans | calls=1 | time= 262.104 ms/call // ===== summary end ===== // -// Those first two lines are one graph, before and after cuDNN was asked to support it, which is why -// a CREATE_GRAPH with no CACHE_GRAPH following it is a refusal -- the one event pattern that says a -// build was paid for and thrown away. -// -// Rows for a site a thread never reached are left out rather than zeroed, which is -// why tid=1 has a backward row and no forward one: in a PyTorch step the forward and -// the backward's support probe run on the main thread, and the backward itself on the -// autograd thread, which finds the graph that probe left behind. That split is why -// the build identities hold on the totals rows and not on any single thread's. -// -// Reading this file: the interface is the four names under "vocabulary" and the six recorders at -// the bottom, and that is everything the rest of the library touches. In between, in namespace -// detail, is what they are built out of, in the order an event travels through it -- the gate, the -// counters, the line, the exit summary. A question about what the output means is answered by the -// counter definitions in the middle; a question about what to call is answered by the bottom. -// ============================================================================ +// Those first two lines are one graph before and after cuDNN was asked to support it, which is why +// a CREATE_GRAPH with no CACHE_GRAPH after it is a refusal. Rows for a site a thread never reached +// are left out rather than zeroed, hence tid=1 having a backward row and no forward one: a PyTorch +// step runs the forward and the backward's probe on the main thread and the backward itself on the +// autograd thread, which finds the graph that probe left behind. That split is why the build +// identities hold on the totals rows and not on any single thread's. #ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ #define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ @@ -78,48 +66,36 @@ namespace fused_attn { namespace graph_cache_debug { // ============================================================================ -// The vocabulary a call site needs: which build site an event came from, and which build stage or -// lookup outcome it is reporting. These four names and the recorders at the bottom of the file are -// the whole interface; everything between them is machinery, in namespace detail. -// -// Backend and Pass are fused_attn's own, from config_and_params.h, so that a recorder and the key -// it prints share one notion of a site; taking the pair rather than the "fwd"/"bwd" strings this -// used to also turns a mistake at a call site into a compile error. -// -// Every recorder names both halves, since the counters are per site -- adding f16's builds into -// fp8's column would leave a run that drove both unable to say which paid for what. +// Vocabulary: which build site an event came from, and which build stage or lookup outcome it +// reports. These four names and the recorders at the bottom are the whole interface. Backend and +// Pass are fused_attn's own, so a recorder and the key it prints share one notion of a site, and a +// mistake at a call site is a compile error rather than a mistyped "fwd". // ============================================================================ inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } -// The frontend calls that make up a build, in the order they run. `kCount` must stay last: it -// sizes the timing table, and detail::kStageNames is indexed by these values when the summary -// prints, so the two must be kept in the same order. +// The frontend calls that make up a build, in the order they run. `kCount` must stay last: it sizes +// the timing table, and detail::kStageNames is indexed by these values, so the two must stay in the +// same order. enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; // What a lookup found: an entry, or nothing. enum class LookupResult { Miss, Hit }; // ============================================================================ -// Machinery: the gate, the counters, the formatting and the exit summary. Nothing outside this -// file names any of it. -// -// Reading order below is the order an event travels: whether to record at all, which site it -// belongs to, the counters it moves, the line it prints, and finally the summary that reports the -// lot at exit. +// Machinery, in the order an event travels through it: the gate, the site index, the counters, the +// line, the exit summary. Nothing outside this file names any of it. // ============================================================================ namespace detail { // ============================================================================ // The gate: whether this process records anything, and how it names itself when it does. Every -// answer here is fixed for the life of the process and read out of an initialized-once static, so -// the check a disabled build pays at each call site is one load and one branch. +// answer here is fixed for the life of the process and cached in an initialized-once static, so a +// disabled build pays one load and one branch per call site. // ============================================================================ -// Verbosity level parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG (0=off, 1=events, -// 2=trace). Single read at startup, cached; when unset every call site pays one -// cached-flag check and nothing else. +// Verbosity parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG: 0=off, 1=events, 2=trace. inline int debug_level() { static const int lvl = [] { const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); @@ -130,8 +106,8 @@ inline int debug_level() { return lvl; } -// Rank of this process as reported by the launcher, or -1 when there is no -// launcher (a single-process run). First variable that is set wins. +// Rank of this process as reported by the launcher, or -1 when there is none. First variable that +// is set wins. inline int launcher_rank() { static const int rank = []() -> int { for (const char *var : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { @@ -143,15 +119,11 @@ inline int launcher_rank() { return rank; } -// Diagnostics are on at level >= 1, and only for the ranks the ":" suffix selects. Every -// rank writes to the same stderr, and under data/tensor parallelism they run identical shapes, so -// emitting from all of them multiplies the volume by the world size to say the same thing. Hence -// rank 0 only by default. Context parallelism is the case worth overriding for: the ranks run -// different subsets of the per-step regimes, so their build counts genuinely differ. -// -// Both inputs are fixed for the life of the process, so the whole verdict is one initialized-once -// static -- which is all the check every call site makes reads, the per-lookup path included. -// Unselected ranks skip the counters too, so they pay nothing beyond it. +// On at level >= 1, and only for the ranks the ":" suffix selects. Every rank writes to the +// same stderr and under data/tensor parallelism they run identical shapes, so emitting from all of +// them multiplies the volume by the world size to say the same thing; hence rank 0 only by default. +// Context parallelism is the case worth overriding for, its ranks running different subsets of the +// per-step regimes. inline bool enabled() { static const bool on = [] { if (debug_level() < 1) return false; @@ -175,16 +147,14 @@ inline bool enabled() { return on; } -// The gate on the per-lookup and per-execution trace lines: everything enabled() asks for, level 2 -// on top of it. Named for that conjunction, and testing it rather than just the level, so the -// answer holds wherever it is asked -- level 2 alone is true on a rank that emits nothing, which -// would make this read as "trace" on every rank in the job. +// The gate on the per-lookup and per-execution trace lines. Tests enabled() rather than just the +// level, so the answer holds wherever it is asked: level 2 alone is true on a rank that emits +// nothing, which would make this read as "trace" on every rank in the job. inline bool enabled_with_trace() { return enabled() && debug_level() >= 2; } // Names the emitting rank, without which the ranks sharing one stderr would be indistinguishable. -// A run whose launcher exports no rank is left untagged rather than falling back to a pid, an -// OS-level identifier only being useful for correlating against a profiler. The tag carries its -// own trailing separator, so the untagged case prints no empty column. +// The tag carries its own trailing separator, so a run whose launcher exports no rank prints no +// empty column. inline const std::string &rank_tag() { static const std::string *tag = [] { const int rank = launcher_rank(); @@ -195,97 +165,59 @@ inline const std::string &rank_tag() { } // Short thread IDs (0, 1, 2, ...) in assignment order, not identity: tid=0 is whichever thread -// touched this cache first, and the number means nothing outside this process. It attributes the -// per-thread summary rows and is not meant to be matched against anything external. +// touched this cache first, and the number means nothing outside this process. 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. +// Registered at first use. On process exit, prints event counters and build timings. inline void register_summary_once(); -// ============================================================================ -// Indexing the build site an event came from: f16 or fp8, forward or backward. -// ============================================================================ - -// Backend major, pass minor, so that the two passes of one backend are adjacent -- which is how -// the counter lines and the summary rows present them, one backend at a time. +// Backend major, pass minor, so that the two passes of one backend are adjacent -- which is how the +// counter lines and the summary rows present them. constexpr size_t kSiteCount = 4; inline constexpr size_t site_index(Backend b, Pass p) { return (b == Backend::F16 ? 0u : 2u) + (p == Pass::Fwd ? 0u : 1u); } // ============================================================================ -// Cache event counters, one block per build site. Each name is both the event tag on the line -// that records it and the column carrying its running total: -// - create_graph: a graph constructed for a miss, counted before cuDNN is asked to support it and -// so regardless of what cuDNN goes on to say about it. +// Cache event counters, one block per build site. Each name is both the event tag on the line that +// records it and the column carrying its running total: +// - create_graph: a graph constructed for a miss, counted before cuDNN is asked about it. // - cache_graph: one of those graphs cleared check_support(), so this is the graphs cuDNN agreed -// to run. Counted on that verdict rather than on the insert that follows it, so it says what -// cuDNN accepted and not how many entries the map holds; the two differ only when a build race -// is lost and a supported graph is discarded for the winner's. -// - build_plans: a cached graph finished with graph.build_plans(), the kernel compilation that -// cache_graph deferred. At most one per cache_graph, paid by that graph's first execution -// rather than by the probe that built it. -// - execute: a graph execution cuDNN accepted, counted once the enqueue returns. Not a completed +// to run. Counted on that verdict rather than on the insert that follows, so it says what cuDNN +// accepted and not how many entries the map holds. +// - build_plans: a cached graph finished with graph.build_plans(), the kernel compilation +// cache_graph deferred, paid by that graph's first execution rather than by the probe. +// - execute: an execution cuDNN accepted, counted once the enqueue returns. Not a completed // execution: the work is asynchronous, so a device-side fault is not reflected here. // - hit: a lookup answered from the cache. Need not lead to an execution -- it can be a backend -// availability check, or the workspace-sizing call of nvte_fused_attn_fwd/bwd, which has no -// runtime tensors to run with. -// - miss: a lookup the cache did not answer; triggers a graph build. +// availability check, or a workspace-sizing call, which has no tensors to run with. +// - miss: a lookup the cache did not answer; triggers a build. // -// Identities, holding by construction, so a violation is a bug in the cache or in the counting -// rather than something the workload did: -// - hit + miss = every lookup, one per call to cache_graph(), which makes it the denominator for -// everything below. (The function, not the column of the same name -- the column counts the -// subset of those calls that ended in an entry.) -// - miss >= create_graph >= cache_graph, where each drop is a build that threw. -// create_graph - cache_graph is what cuDNN refused, or could not reach a verdict on; nothing is -// cached for it, so this gap is where a refusal shows up, and the reason itself goes to the -// framework. miss - create_graph would be a backend refusing a configuration from inside its -// own build, and no backend does that any more -- TE's own FP8 rules over bias, ALiBi and the -// recipes it does not serve moved to nvte_get_fused_attn_backend_v2, which answers them with a -// reason rather than throwing. So this gap should read zero, and the column stays as the thing -// that says so: a nonzero miss - create_graph means a build threw where none is expected to. +// Identities, holding by construction, so a violation is a bug in the cache or in the counting: +// - hit + miss = every lookup, one per call to cache_graph() (the function, not the column of the +// same name), which makes it the denominator for the rest. +// - miss >= create_graph >= cache_graph, each drop a build that threw. create_graph - +// cache_graph is what cuDNN refused, and is where a refusal shows up. miss - create_graph would +// be a backend refusing from inside its own build, which none does since TE's FP8 rules moved +// to nvte_get_fused_attn_backend_v2 to be answered with a reason instead, so that gap should +// read zero and the column stays as the thing that says so. // - cache_graph >= build_plans, the gap being graphs a probe built that nothing has run. -// - execute > 0 implies build_plans > 0, every site calling build_plans() ahead of the -// workspace-sizing return, itself ahead of record_execute. Read backwards: a workspace-sizing -// call pays build_plans and never execute. +// - execute > 0 implies build_plans > 0, every site building plans ahead of the workspace-sizing +// return, itself ahead of record_execute. So a sizing call pays build_plans and never execute. // - Both build identities belong to the totals rows, not to one thread's: the thread that builds // a graph need not compile its plans, and a PyTorch step splits exactly that way. // - Per-thread rows sum column by column to "tid=all dev=all", and the per-backend rows of one // pass to that pass's all-backends row. -// - A lost build race disturbs none of the above: the loser records its own miss, create_graph -// and cache_graph, having built a graph cuDNN did agree to run, and the once_flag still permits -// one build_plans on the winner's entry. What it costs is a build, which two MISS lines on one -// key is the way to see. // - Stage timing calls fall along validate >= build_operation_graph >= create_execution_plans >= -// check_support, each drop being the builds that ended at the stage before, which localizes -// where cuDNN refuses rather than only how long refusing took. -// - The build_plans timing row can show more calls than the build_plans column, the difference -// being plan builds that threw: the timer records while unwinding, the counter only on return. +// check_support, each drop being builds that ended at the stage before, which localizes where +// cuDNN refuses rather than only how long refusing took. The build_plans timing row can exceed +// its column: the timer records while unwinding, the counter only on return. // -// Signatures, workload-dependent, so read rather than asserted. What a column stalling says about -// who rejected a configuration is the user-facing half of this and lives in docs/envvars.rst; what -// follows is what is worth knowing on top of it: -// - After warmup only hit and execute should move; a late create_graph means something varies per -// step that need not. -// - Several hits per execution is normal, since selection, workspace sizing and execution all -// look the same key up; what matters is that the ratio stays flat. -// - execute / cache_graph is the amortization figure, and a lower bound at that, a lost race -// counting a supported graph the cache did not keep. Single digits after a long run means the -// cache is not earning its keep. -// - miss climbing without settling means the key space is not closing, and since the cache is -// unbounded, every distinct key is held for the life of the process. -// - A build count that looks doubled on a multi-device process usually is not: device_id is part -// of the key, so the same shape on two devices is two entries. Read the dev column. -// - Two MISS lines with the same key, cache_graph above the number of distinct keys, is that lost -// race: wasted work rather than a bug, worth chasing only if it repeats. -// - A level-2 trace is the set of lookups, not their order, the line being written after the -// cache lock is dropped. +// What the columns say about a workload is user-facing and lives in docs/envvars.rst. // ============================================================================ struct EventCounters { @@ -303,9 +235,8 @@ inline EventCounters &counters(Backend b, Pass p) { } // One counter block read out into plain values, so the summary can sum blocks for its per-backend -// and all-backends rows. The columns are not read as one indivisible operation, which nothing here -// wants: the summary runs at exit, after the writing threads are done, and an event line is a -// snapshot of a moving count by nature. +// and all-backends rows. Not read as one indivisible operation, which nothing here wants: the +// summary runs at exit, after the writing threads are done. struct CounterSnapshot { uint64_t create_graph = 0; uint64_t cache_graph = 0; @@ -324,8 +255,8 @@ struct CounterSnapshot { return *this; } - // Whether this block saw nothing at all, which is what lets the summary leave out the rows - // for a backend the run never used rather than printing zeros for it. + // Whether this block saw nothing at all, which is what lets the summary leave out the rows for a + // backend the run never used rather than printing zeros for it. bool empty() const { return (create_graph | cache_graph | build_plans | execute | hit | miss) == 0; } @@ -342,17 +273,13 @@ inline CounterSnapshot snapshot(const EventCounters &c) { return s; } -// ============================================================================ -// The same counters again, per thread, and the registry the exit summary walks to find them. -// ============================================================================ - -// Per-thread counters, one block per build site, so the summary can break every column down by -// thread and backend: in the single-process context-parallel case each device is driven by its own -// thread, and under PyTorch this separates the main thread from the autograd one. +// Per-thread counters, so the summary can break every column down by thread and backend: in the +// single-process context-parallel case each device is driven by its own thread, and under PyTorch +// this separates the main thread from the autograd one. // // `device` is the device this thread last drove, restamped on every event. Event lines print the // live current device instead, which is exact; this exists for the per-thread summary rows, written -// at exit by whichever thread is exiting, which cannot ask the recorded thread what it was doing. +// at exit by whichever thread is exiting. struct ThreadCounters { unsigned tid = 0; std::atomic device{-1}; @@ -361,9 +288,9 @@ struct ThreadCounters { // The registry and its mutex are heap-allocated and deliberately never freed. Static destructors // and atexit handlers run as one sequence in reverse order of construction, and this registry is -// built lazily, so it can be constructed *after* the summary handler is registered -- and would -// then be destroyed *before* it runs, leaving the handler to lock a destroyed mutex and walk a -// destroyed vector. Leaking removes the ordering question, at a cost of one mutex and one vector. +// built lazily, so it can be constructed after the summary handler is registered -- and would then +// be destroyed before it runs, leaving the handler to lock a destroyed mutex and walk a destroyed +// vector. Leaking removes the ordering question, at a cost of one mutex and one vector. inline std::mutex &thread_registry_mutex() { static std::mutex *m = new std::mutex(); return *m; @@ -373,9 +300,8 @@ inline std::vector &thread_registry() { return *v; } -// This thread's counter block, leaked for a related but distinct reason: a worker thread can exit -// long before the process does, while the registry holds a pointer to its block for the exit -// summary. Tying the block's lifetime to the thread would leave that pointer dangling. +// This thread's block, leaked for a related but distinct reason: a worker thread can exit long +// before the process does, while the registry holds a pointer to its block for the exit summary. inline ThreadCounters &thread_counters() { static thread_local ThreadCounters *tc = [] { auto *p = new ThreadCounters(); @@ -398,20 +324,15 @@ inline EventCounters &thread_counters(Backend b, Pass p) { // ============================================================================ // Turning a counter block into a line, and getting a line out. One formatter, shared by the event -// lines and the summary rows, so that the two cannot drift into presenting the same columns -// differently, and one writer, so that everything here reaches stderr the same way. +// lines and the summary rows so the two cannot drift apart, and one writer, so everything here +// reaches stderr the same way. // ============================================================================ // The one place diagnostics reach stderr, and the reason it exists: the first line this process // writes carries a leading newline. Diagnostics share stderr with whatever the framework is -// printing, and a test runner's progress output has no trailing newline of its own, so without -// this the first line continues someone else's -- which on a level-2 trace line, long enough to -// wrap already, leaves no way to find where it starts. Where the previous output did end cleanly -// the prefix reads as a blank line setting the diagnostics apart from it. -// -// One fwrite per line either way: a rank's summary block is assembled whole precisely so that -// concurrently exiting ranks do not interleave, and the extra allocation buys the same for the one -// line that gets the prefix. +// printing, and a test runner's progress output has no trailing newline of its own, so without this +// the first line continues someone else's. Where the previous output did end cleanly the prefix +// reads as a blank line setting the diagnostics apart. inline void write_stderr(const std::string &text) { static std::atomic first_line{true}; if (first_line.exchange(false, std::memory_order_relaxed)) { @@ -424,21 +345,14 @@ inline void write_stderr(const std::string &text) { } // Format one counter block -- one pass of one backend -- as one line. One pass rather than both -// because a line carrying the forward and backward columns together ran past 300 characters and -// wrapped in most terminals; the two passes are adjacent rows instead. -// -// `tid_field` and `dev_field` are whole columns, e.g. "tid=3" and "dev=0". The totals rows pass -// "tid=all" and "dev=all", since those counters are summed across whatever the process drove and -// naming one thread or device would be a lie. +// because a line carrying the forward and backward columns together ran past 300 characters; the +// two passes are adjacent rows instead. // -// `label` is the build site, "f16 fwd", plus the event name on an event line, and arrives padded to -// the width its own kind of line uses: 20 characters for an event line, 7 for a summary row. -// Deliberately not one width for both -- sharing it would put twelve blank columns on every summary -// row to align the scattered event lines against a block that is delimited and read on its own. -// -// The thread and device come first, so every line, level-2 trace lines included, shares one prefix -// to read down. What the columns mean and the identities they satisfy are with the definitions -// above. +// `tid_field` and `dev_field` are whole columns, e.g. "tid=3" and "dev=0"; the totals rows pass +// "tid=all" and "dev=all". `label` is the build site, "f16 fwd", plus the event name on an event +// line, and arrives padded to the width its own kind of line uses -- 20 characters for an event +// line, 7 for a summary row -- deliberately not one width for both, which would put twelve blank +// columns on every summary row to align the scattered event lines. inline std::string format_counter_line(const char *tid_field, const char *dev_field, const char *label, const CounterSnapshot &c) { char buf[512]; @@ -453,8 +367,7 @@ inline std::string format_counter_line(const char *tid_field, const char *dev_fi // One event line, from the thread the event happened on, carrying the running totals of the build // site that raised it. The device is read live rather than remembered, so it is the device this -// event was actually issued against, and is recorded on the thread's block on the way past for -// the benefit of the exit summary. +// event was actually issued against, and is recorded on the thread's block for the exit summary. inline void print_counters(Backend b, Pass p, const char *event) { const int device = cuda::current_device(); thread_counters().device.store(device, std::memory_order_relaxed); @@ -469,18 +382,13 @@ inline void print_counters(Backend b, Pass p, const char *event) { write_stderr(format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p)))); } -// ============================================================================ -// What the recorders at the bottom of the file are made of: moving one column, and naming a -// lookup's outcome. -// ============================================================================ - -// The body every recorder shares: gate, register the exit summary, and add one to `column` -// in both the process-wide block and this thread's. Returns whether diagnostics are on at all, so -// that a caller can skip building a line nobody will read. +// The body every recorder shares: gate, register the exit summary, and add one to `column` in both +// the process-wide block and this thread's. Returns whether diagnostics are on at all, so a caller +// can skip building a line nobody will read. // -// Both blocks or neither. A recorder that moved one and not the other would leave the per-thread -// rows failing to add up to the totals row, which the summary presents as an invariant, and the -// discrepancy would look like a threading bug in the cache rather than a miscount here. +// Both blocks or neither. Moving one and not the other would leave the per-thread rows failing to +// add up to the totals row, which the summary presents as an invariant, and the discrepancy would +// look like a threading bug in the cache rather than a miscount here. inline bool record_counter(Backend b, Pass p, std::atomic EventCounters::*column) { if (!enabled()) return false; register_summary_once(); @@ -489,9 +397,8 @@ inline bool record_counter(Backend b, Pass p, std::atomic EventCounter return true; } -// The column a lookup lands in, and the tag naming it. Both are written as a switch with no -// default so that adding an outcome fails to compile here rather than being silently counted as -// a miss. +// The column a lookup lands in, and the tag naming it. Both are switches with no default, so adding +// an outcome fails to compile here rather than being silently counted as a miss. inline std::atomic EventCounters::*lookup_column(LookupResult result) { switch (result) { case LookupResult::Hit: @@ -513,39 +420,31 @@ inline const char *lookup_name(LookupResult result) { } // ============================================================================ -// Graph build timings. -// -// Which stage dominates determines what to do about a slow build: time in -// `check_support` and `build_plans` is heuristic selection and kernel compilation, -// largely intrinsic to the shape, while time in `validate` or -// `build_operation_graph` is graph-construction cost on our side. One duration per -// build cannot make that distinction. +// Graph build timings. Which stage dominates says what to do about a slow build: time in +// check_support and build_plans is heuristic selection and kernel compilation, largely intrinsic to +// the shape, while time in validate or build_operation_graph is graph-construction cost on our +// side. One duration per build could not make that distinction. // -// Each stage is wrapped where it is called, in graph_cache.h, and accumulates into -// the table below under its build site. Only sums are kept, so the summary can -// report a mean and nothing else -- and since a build happens once per distinct -// cache key, those calls span different shapes rather than repeating one. Read a +// Only sums are kept, so the summary reports a mean and nothing else -- and since a build happens +// once per distinct cache key, those calls span different shapes rather than repeating one. Read a // stage mean as where build time goes in aggregate, not as any one build's cost. // ============================================================================ -// Indexed by BuildStage when the summary prints, so it must stay in that enum's order and carry -// one name per stage ahead of its kCount sentinel. +// Indexed by BuildStage when the summary prints, so it must stay in that enum's order and carry one +// name per stage ahead of its kCount sentinel. inline constexpr const char *kStageNames[] = { "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; -// Totals for one (pass, stage) pair. Relaxed ordering is sufficient: these -// counters order nothing, and the only read happens once, after the threads that -// wrote them are done. +// Totals for one (pass, stage) pair. Relaxed ordering is sufficient: these counters order nothing, +// and the only read happens once, after the threads that wrote them are done. struct StageTiming { std::atomic calls{0}; std::atomic time_ns{0}; }; -// Bucketed by build site, so the summary can report the cost of each stage separately for each -// backend and pass -- an fp8 build and an f16 build are different work, and averaging them -// together would describe neither. Unlike the thread registry above, this table needs no leak to -// outlive the exit handler that reads it: it holds nothing but atomics, so it is trivially -// destructible and no destructor is registered for it at all. +// Bucketed by build site, an fp8 build and an f16 build being different work. Unlike the thread +// registry above, this needs no leak to outlive the exit handler that reads it: it holds nothing +// but atomics, so no destructor is registered for it at all. constexpr size_t kStageBuckets = kSiteCount * static_cast(BuildStage::kCount); inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { static std::array table{}; @@ -554,12 +453,10 @@ inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { return table[idx]; } -// Times one stage: clock read in the constructor, accumulated in the destructor. -// Recording on scope exit rather than at an explicit stop() keeps a failing stage -// measurable, since `build_plans` throws through NVTE_CHECK_CUDNN_FE and the -// destructor still runs while unwinding, so a build that dies there contributes its -// time instead of vanishing. `on` is latched at construction rather than re-tested in -// the destructor, so the destructor can never accumulate against an unset `start`. +// Times one stage: clock read in the constructor, accumulated in the destructor. Recording on scope +// exit rather than at an explicit stop() keeps a failing stage measurable, since build_plans throws +// through NVTE_CHECK_CUDNN_FE and the destructor still runs while unwinding. `on` is latched at +// construction, so the destructor can never accumulate against an unset `start`. struct ScopedBuildTimer { BuildStage stage; bool on; @@ -584,13 +481,8 @@ struct ScopedBuildTimer { }; // ============================================================================ -// Summary: on process exit, print cache event counters and graph build timings. -// -// Each section below appends its rows to the block the handler is assembling, in the order they -// are printed: per-thread rows, then totals, then stage timings. Split into named pieces rather -// than written inline because they are read one at a time -- a question about the output is a -// question about one of these -- and because the registration itself is already three constructs -// deep (an initialized-once static holding an atexit handler) before any row logic joins it. +// Summary: on process exit, print cache event counters and graph build timings. Each piece below +// appends its rows to the block the handler is assembling, in the order they are printed. // ============================================================================ // The two backends that keep a cache, in the order every part of the summary walks them. @@ -637,10 +529,10 @@ inline void append_thread_rows(std::string &block) { } } -// Totals, printed after the per-thread rows so they read as their sum: one row per build site, -// then one per pass across the backends when the run used more than one. Both come from the -// process-wide counters rather than by adding up the rows above, so the two agreeing is a check -// on the counting rather than an artifact of it. +// Totals, printed after the per-thread rows so they read as their sum: one row per build site, then +// one per pass across the backends when the run used more than one. Both come from the process-wide +// counters rather than by adding up the rows above, so the two agreeing is a check on the counting +// rather than an artifact of it. inline void append_total_rows(std::string &block) { CounterSnapshot all_fwd; CounterSnapshot all_bwd; @@ -661,8 +553,7 @@ inline void append_total_rows(std::string &block) { } } -// Mean time per call for each stage of each build site, skipping stages nothing reached. A mean is -// all the sums kept can support; see the section above for why that is the right figure to read. +// Mean time per call for each stage of each build site, skipping stages nothing reached. inline void append_stage_rows(std::string &block) { for (const Backend b : kSummaryBackends) { for (const Pass p : {Pass::Fwd, Pass::Bwd}) { @@ -687,8 +578,8 @@ inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; - // Built in memory and emitted with one write, so that concurrently-exiting - // processes (one per rank under torchrun) stay grouped rather than interleaving. + // Built in memory and emitted with one write, so that concurrently-exiting processes (one per + // rank under torchrun) stay grouped rather than interleaving. const std::string marker = "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary "; std::string block = marker + "begin =====\n"; append_thread_rows(block); @@ -710,21 +601,13 @@ inline void register_summary_once() { // // Every one of them is called after the event it names, never before, so that a column counts what // happened rather than what was attempted. That is what gives the gaps between columns their -// meaning: an event that can fail partway -- a build cuDNN refuses, an execution whose setup throws -// first -- leaves the earlier column moved and the later one not. -// -// record_time is the exception, and only because timing cannot be done after the fact: it brackets -// the call it measures, and accumulates while unwinding so that a stage which throws is still -// timed. Its timing rows can therefore outnumber the matching counter column. -// -// Which of them belongs where in the cache's flow is documented on each below and in graph_cache.h -// at the call sites. +// meaning: an event that fails partway leaves the earlier column moved and the later one not. +// record_time is the exception, and only because timing cannot be done after the fact. // ============================================================================ -// A graph constructed for a miss, whatever cuDNN goes on to make of it. Call from the miss path -// that built it, as soon as construction returns and before check_support() is asked. Before, -// because construction is where a backend would refuse a configuration on its own rules, and such -// a build never gets here -- which is what makes miss - create_graph builds that failed on this +// A graph constructed for a miss, whatever cuDNN goes on to make of it. Call as soon as +// construction returns and before check_support() is asked: construction is where a backend would +// refuse on its own rules, which is what makes miss - create_graph the builds that failed on this // side of cuDNN. No backend does that now, so the gap is there to read as zero. inline void record_create_graph(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::create_graph)) { @@ -732,36 +615,28 @@ inline void record_create_graph(Backend b, Pass p) { } } -// A created graph that cleared check_support(), so this counts the graphs cuDNN agreed to run. Call -// as soon as that verdict returns, ahead of the insert: a refused graph throws in between, leaving -// its CREATE_GRAPH unanswered, which is what makes create_graph - cache_graph cuDNN's refusals. -// Deliberately not the insert, so that a lost race reads as the extra build it is rather than as a -// count that disagrees with the size of the cache. +// A created graph that cleared check_support(). Call as soon as that verdict returns, ahead of the +// insert: a refused graph throws in between, leaving its CREATE_GRAPH unanswered, which is what +// makes create_graph - cache_graph cuDNN's refusals. inline void record_cache_graph(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::cache_graph)) { detail::print_counters(b, p, "CACHE_GRAPH"); } } -// The graph.build_plans() a cache_graph deferred, now completed. Call from inside the -// std::call_once that runs it, and after the call returns rather than before: it throws without -// setting the once_flag, leaving a later execution to retry, so counting on the way out keeps this -// a count of graphs that reached a runnable state. +// The graph.build_plans() a cache_graph deferred, now completed. Call from inside the call_once +// that runs it, and after the call returns: it throws without setting the once_flag, leaving a +// later execution to retry, so counting on the way out keeps this a count of runnable graphs. inline void record_build_plans(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::build_plans)) { detail::print_counters(b, p, "BUILD_PLANS"); } } -// An execution cuDNN accepted. Call after graph.execute() returns, as with the recorders above, so -// that a graph the surrounding setup never reached is not counted as having run -- the stream set -// and the cu_seqlens conversion kernels sit between the decision to execute and the execution, and -// either can throw. -// -// Accepted is as far as this can go. execute() enqueues on a stream and returns, so a fault the -// device raises later, surfacing at the next synchronization, still leaves the execution counted -// here. There is no synchronous completion point to hook without making the diagnostic change what -// it measures. +// An execution cuDNN accepted. Call after graph.execute() returns, so a graph the surrounding +// setup never reached is not counted as having run -- the stream set and the cu_seqlens conversion +// kernels sit in between, and either can throw. Accepted is as far as this can go: execute() +// enqueues and returns, so a fault the device raises later still leaves the execution counted here. // // Unlike the recorders above, this fires on every execution rather than once per distinct key, so // its line is held back to level 2 while its column keeps counting. @@ -774,19 +649,13 @@ inline void record_execute(Backend b, Pass p) { // `key` is the normalized cache key -- make_cache_key(pass)'s output, the exact value looked up -- // not the execution config it came from. HIT/MISS is decided by comparing keys, so a trace of -// anything else cannot explain its own outcome: the pre-normalization config would show identical -// lines with opposite outcomes, and differing lines that both hit. Diffing two MISS lines here -// names exactly the fields responsible for the extra build. -// -// The cost is that overwritten fields are no longer visible in their original form: attn_scale -// reads 1, ragged num_tokens read 0, max_seqlen and batch_size read their bucketed values. -// -// This is the one line here not built from counters, so it does not go through -// format_counter_line: which fields it names is FusedAttnConfig::key_debug_string()'s to say, -// alongside the operator< that decides what a key compares on in the first place. +// anything else could not explain its own outcome, and diffing two MISS lines here names exactly +// the fields responsible for the extra build. The cost is that overwritten fields no longer appear +// in their original form: attn_scale reads 1, ragged num_tokens read 0, max_seqlen and batch_size +// read their bucketed values. inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedAttnConfig &key) { - // The per-lookup config dump is the highest-volume line (one per cache lookup); - // keep it out of the level-1 path and off the stderr lock unless tracing. + // The highest-volume line here, one per cache lookup; keep it out of the level-1 path and off the + // stderr lock unless tracing. if (!detail::record_counter(b, p, detail::lookup_column(result)) || !detail::enabled_with_trace()) { return; @@ -796,13 +665,15 @@ inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedA "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | ", detail::rank_tag().c_str(), detail::thread_seq_id(), key.device_id, backend_name(b), pass_name(p), detail::lookup_name(result)); - detail::write_stderr(prefix + key.key_debug_string() + "\n"); + // The one line here not built from counters: which fields it names is + // FusedAttnConfig::to_string()'s to say, alongside the operator< that decides what a key + // compares on in the first place. + detail::write_stderr(prefix + key.to_string() + "\n"); } // Record how long `fn` takes as `stage` of the given build site. Unlike the recorders above this -// wraps the work rather than reporting on work already done, which is the point: the measured -// region is exactly the call passed in, so surrounding work cannot drift into it as that code -// changes. Stage timings feed the summary only; they print no line of their own. +// wraps the work rather than reporting on work already done, so the measured region is exactly the +// call passed in. Stage timings feed the summary only; they print no line of their own. // // Passes `fn`'s result back out so that a timed call reporting a value can be written as the // initializer of that value. cuDNN's error_t is [[nodiscard]], and the alternative -- declaring the diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index d7106803d0..691881dcb2 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -136,9 +136,10 @@ enum NVTE_Mask_Type { * \brief Attention softmax types as described in * Efficient Streaming Language Models with Attention Sinks (https://arxiv.org/pdf/2309.17453v3). * For a given attention score S = Q*K^T, different softmax types perform different operations on S, - * NVTE_VANILLA_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - * NVTE_OFF_BY_ONE_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - * NVTE_LEARNABLE_SOFTMAX: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), + * `NVTE_VANILLA_SOFTMAX`: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), + * `NVTE_OFF_BY_ONE_SOFTMAX`: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and + * `NVTE_LEARNABLE_SOFTMAX`: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + + * sum(exp(S[:,j,:,:]), dim=-1)), * where alpha is a learnable parameter of shape [H]. */ enum NVTE_Softmax_Type { @@ -198,11 +199,11 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); typedef void *NVTEFusedAttnConfig; /*! \enum NVTEFusedAttnConfigAttribute - * \brief Attributes 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 ``attr_sizes`` array of ``FusedAttnConfig``. New fields may only be appended - * at the end and 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 `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 settings @@ -277,11 +278,11 @@ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, typedef void *NVTEFusedAttnFwdParams; /*! \enum NVTEFusedAttnFwdParamsAttribute - * \brief Attributes 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. + * 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 @@ -290,6 +291,9 @@ enum NVTEFusedAttnFwdParamsAttribute { kNVTEFusedAttnFwdParamsV, kNVTEFusedAttnFwdParamsBias, kNVTEFusedAttnFwdParamsSoftmaxOffset, + kNVTEFusedAttnFwdParamsS, + kNVTEFusedAttnFwdParamsO, + kNVTEFusedAttnFwdParamsAuxCtxTensors, kNVTEFusedAttnFwdParamsCuSeqlensQ, kNVTEFusedAttnFwdParamsCuSeqlensKV, kNVTEFusedAttnFwdParamsCuSeqlensQPadded, @@ -297,26 +301,23 @@ enum NVTEFusedAttnFwdParamsAttribute { kNVTEFusedAttnFwdParamsPageTableK, kNVTEFusedAttnFwdParamsPageTableV, kNVTEFusedAttnFwdParamsRngState, - kNVTEFusedAttnFwdParamsS, - kNVTEFusedAttnFwdParamsO, - kNVTEFusedAttnFwdParamsAuxCtxTensors, // configuration knobs + kNVTEFusedAttnFwdParamsMaxSeqlenQ, + kNVTEFusedAttnFwdParamsMaxSeqlenKV, kNVTEFusedAttnFwdParamsIsTraining, - kNVTEFusedAttnFwdParamsCudaGraph, kNVTEFusedAttnFwdParamsReturnMaxLogit, - kNVTEFusedAttnFwdParamsAttnMaskType, + kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsAttnScale, + kNVTEFusedAttnFwdParamsDropout, + kNVTEFusedAttnFwdParamsQKVLayout, + kNVTEFusedAttnFwdParamsOFormat, + kNVTEFusedAttnFwdParamsQKVScaleInvFormat, kNVTEFusedAttnFwdParamsBiasType, + kNVTEFusedAttnFwdParamsAttnMaskType, kNVTEFusedAttnFwdParamsSoftmaxType, kNVTEFusedAttnFwdParamsWindowSizeLeft, kNVTEFusedAttnFwdParamsWindowSizeRight, kNVTEFusedAttnFwdParamsBottomRightDiagonal, - kNVTEFusedAttnFwdParamsDropout, - kNVTEFusedAttnFwdParamsAttnScale, - kNVTEFusedAttnFwdParamsQKVLayout, - kNVTEFusedAttnFwdParamsOFormat, - kNVTEFusedAttnFwdParamsQKVScaleInvFormat, - kNVTEFusedAttnFwdParamsMaxSeqlenQ, - kNVTEFusedAttnFwdParamsMaxSeqlenKV, // workspace and stream kNVTEFusedAttnFwdParamsWorkspace, kNVTEFusedAttnFwdParamsStream, @@ -344,11 +345,11 @@ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, typedef void *NVTEFusedAttnBwdParams; /*! \enum NVTEFusedAttnBwdParamsAttribute - * \brief Attributes 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. + * 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 @@ -370,24 +371,24 @@ enum NVTEFusedAttnBwdParamsAttribute { kNVTEFusedAttnBwdParamsCuSeqlensQPadded, kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, // configuration knobs - kNVTEFusedAttnBwdParamsCudaGraph, - kNVTEFusedAttnBwdParamsDeterministic, - kNVTEFusedAttnBwdParamsAttnMaskType, - kNVTEFusedAttnBwdParamsBiasType, - kNVTEFusedAttnBwdParamsSoftmaxType, - kNVTEFusedAttnBwdParamsWindowSizeLeft, - kNVTEFusedAttnBwdParamsWindowSizeRight, - kNVTEFusedAttnBwdParamsBottomRightDiagonal, - kNVTEFusedAttnBwdParamsDropout, + kNVTEFusedAttnBwdParamsMaxSeqlenQ, + kNVTEFusedAttnBwdParamsMaxSeqlenKV, kNVTEFusedAttnBwdParamsAttnScale, + kNVTEFusedAttnBwdParamsDropout, kNVTEFusedAttnBwdParamsQKVLayout, kNVTEFusedAttnBwdParamsOFormat, kNVTEFusedAttnBwdParamsDOFormat, kNVTEFusedAttnBwdParamsDQKVLayout, kNVTEFusedAttnBwdParamsQKVScaleInvFormat, kNVTEFusedAttnBwdParamsDOScaleInvFormat, - kNVTEFusedAttnBwdParamsMaxSeqlenQ, - kNVTEFusedAttnBwdParamsMaxSeqlenKV, + kNVTEFusedAttnBwdParamsBiasType, + kNVTEFusedAttnBwdParamsAttnMaskType, + kNVTEFusedAttnBwdParamsSoftmaxType, + kNVTEFusedAttnBwdParamsWindowSizeLeft, + kNVTEFusedAttnBwdParamsWindowSizeRight, + kNVTEFusedAttnBwdParamsBottomRightDiagonal, + kNVTEFusedAttnBwdParamsDeterministic, + kNVTEFusedAttnBwdParamsCudaGraph, // workspace and stream kNVTEFusedAttnBwdParamsWorkspace, kNVTEFusedAttnBwdParamsStream, @@ -414,21 +415,21 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, /*! \brief Get fused-attention backend based on user configuration. * * This function passes the user configuration to cuDNN frontend, runs its support checks, - * and returns a backend if supported, otherwise a message explaining why the configuration is not supported. - * If supported, the backend is cached and reused for future calls with the same configuration. + * and returns a backend if supported, or `NVTE_No_Backend` and a message explaining why. + * If the configuration is supported, the backend is cached and reused for future calls. * * \param[in] cfg Fused-attention configuration created by - * ``nvte_create_fused_attn_config()``. - * \param[out] message If the configuration is supported, an empty string; - * 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``, - * if the given configuration is supported; otherwise, ``NVTE_No_Backend``. + * `nvte_create_fused_attn_config()`. + * \param[out] message If the configuration is supported, an empty string. If not supported, + * a diagnostic message explaining why there is no support. Pass `NULL` to + * skip the diagnostics. Also, 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`, + * 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); @@ -455,14 +456,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, * \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. + * \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. + * \note `nvte_get_fused_attn_backend` has a narrower input signature than + * `nvte_get_fused_attn_backend_v2`. It fills the fields that it cannot express with + * `nvte_get_fused_attn_backend_v2`'s default values. This includes setting + * `batch_size` = 1, deriving output/gradient formats from `qkv_layout`, assuming a standard + * bias shape [b, h, sq, skv] for `NVTE_POST_SCALE_BIAS`, using delayed scaling for all FP8, + * and not supporting paged-KV attention. Users who need more precise control should + * switch to `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, @@ -474,9 +476,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( /*! \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 - * destroy it with ``nvte_destroy_fused_attn_fwd_params()``. + * `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. */ @@ -531,7 +533,8 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); * \param[in] qkv_layout QKV tensors' layout. * \param[in] o_format Output format. * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; - * if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout. + * if `NVTE_QKV_Format_NOT_SET`, inferred from + * `qkv_layout`. * \param[in] bias_type Bias type. * \param[in] attn_mask_type Attention mask type. * \param[in] softmax_type Attention softmax type. @@ -541,7 +544,7 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. * - * \deprecated This function has been deprecated in favor of nvte_fused_attn_fwd_v2. + * \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, @@ -561,9 +564,9 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso /*! \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 - * destroy it with ``nvte_destroy_fused_attn_bwd_params()``. + * `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. */ @@ -613,9 +616,11 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); * \param[in] do_format Output gradient's format. * \param[in] dqkv_layout QKV gradient tensors' layout. * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; - * if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout. + * if `NVTE_QKV_Format_NOT_SET`, inferred from + * `qkv_layout`. * \param[in] do_scale_inv_format Format of scale-inverse tensors for dO; - * if NVTE_QKV_Format_NOT_SET, inferred from the output layout. + * if `NVTE_QKV_Format_NOT_SET`, inferred from the + * output layout. * \param[in] bias_type Bias type. * \param[in] attn_mask_type Attention mask type. * \param[in] softmax_type Attention softmax type. @@ -627,7 +632,7 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. * - * \deprecated This function has been deprecated in favor of nvte_fused_attn_bwd_v2. + * \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, @@ -736,7 +741,7 @@ void nvte_cp_thd_read_half_tensor(const NVTETensor &tensor, const NVTETensor &cu * \param[out] lse Output tensor. * \param[in] lse_per_step Input tensor. * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] lse_packed Whether or not lse_per_step is packed. + * \param[in] lse_packed Whether or not `lse_per_step` is packed. * \param[in] stream CUDA stream used for this operation. */ void nvte_cp_thd_second_half_lse_correction(NVTETensor lse, const NVTETensor &lse_per_step, @@ -785,7 +790,7 @@ void nvte_cp_thd_out_correction(NVTETensor out, const NVTETensor &out_per_step, * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. * \param[in] first_half One of ("add", "copy", "none") correction op for first half. * \param[in] second_half One of ("add", "copy", "none") correction op for second half. - Must be different from first_half. + * Must be different from `first_half`. * \param[in] stream CUDA stream used for this operation. */ void nvte_cp_thd_grad_correction(NVTETensor grad, const NVTETensor &grad_per_step, @@ -814,9 +819,9 @@ void nvte_cp_thd_get_partitioned_indices(const NVTETensor &cu_seqlens, NVTETenso * * \param[in] inp Input THD tensor [total_tokens, ...]. * \param[in] cu_seqlens Padded cumulative sequence lengths, [batch_size + 1], int32. - * \param[out] out Output tensor, same shape/dtype as inp. + * \param[out] out Output tensor, same shape/dtype as `inp`. * \param[in] world_size Context-parallel size. - * \param[in] total_tokens Total padded tokens (= inp.shape[0]). + * \param[in] total_tokens Total padded tokens (= `inp.shape[0]`). * \param[in] stream CUDA stream used for this operation. */ void nvte_thd_sequence_order_to_cp_rank_order(const NVTETensor &inp, const NVTETensor &cu_seqlens, @@ -830,9 +835,9 @@ void nvte_thd_sequence_order_to_cp_rank_order(const NVTETensor &inp, const NVTET * * \param[in] inp Input THD tensor [total_tokens, ...]. * \param[in] cu_seqlens Padded cumulative sequence lengths, [batch_size + 1], int32. - * \param[out] out Output tensor, same shape/dtype as inp. + * \param[out] out Output tensor, same shape/dtype as `inp`. * \param[in] world_size Context-parallel size. - * \param[in] total_tokens Total padded tokens (= inp.shape[0]). + * \param[in] total_tokens Total padded tokens (= `inp.shape[0]`). * \param[in] stream CUDA stream used for this operation. */ void nvte_thd_cp_rank_order_to_sequence_order(const NVTETensor &inp, const NVTETensor &cu_seqlens, @@ -847,8 +852,8 @@ void nvte_thd_cp_rank_order_to_sequence_order(const NVTETensor &inp, const NVTET * \param[in] inp Per-split THD source tensor [total_tokens, ...]. * \param[in] cu_seqlens_padded Padded cumulative sequence lengths, [batch_size + 1], int32. * \param[in] cu_seqlens Valid cumulative sequence lengths, [batch_size + 1], int32. - * \param[in,out] out Rank-local accumulator, same shape/dtype as inp. - * \param[in] total_tokens Total padded tokens (= inp.shape[0]). + * \param[in,out] out Rank-local accumulator, same shape/dtype as `inp`. + * \param[in] total_tokens Total padded tokens (= `inp.shape[0]`). * \param[in] stream CUDA stream used for this operation. */ void nvte_thd_copy_valid_tokens_from_per_split_to_rank_local(const NVTETensor &inp, @@ -915,7 +920,7 @@ void nvte_prepare_flash_attn_bwd(NVTETensor q, NVTETensor k, NVTETensor v, NVTET * \param[in] inputs List of input tensors. * \param[in,out] outputs List of output tensors. * \param[in] num_tensors Number of tensors in the list. - * \param[in] original_format Original QKV format (NVTE_BSHD or NVTE_SBHD). + * \param[in] original_format Original QKV format (`NVTE_BSHD` or `NVTE_SBHD`). * \param[in] stream CUDA stream. */ void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs, @@ -983,11 +988,11 @@ class AttentionShape { }; /*! \class FusedAttnConfigWrapper - * \brief C++ helper for constructing an ``NVTEFusedAttnConfig``. + * \brief C++ helper for constructing an `NVTEFusedAttnConfig`. * - * It owns an opaque ``NVTEFusedAttnConfig`` handle created by - * ``nvte_create_fused_attn_config()``, and provides a convenient, - * chainable interface for setting every field in ``FusedAttnConfig``. + * 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: @@ -1160,11 +1165,11 @@ class FusedAttnConfigWrapper { }; /*! \class FusedAttnFwdParamsWrapper - * \brief C++ helper for constructing an ``NVTEFusedAttnFwdParams``. + * \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``. + * 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: @@ -1212,6 +1217,15 @@ class FusedAttnFwdParamsWrapper { FusedAttnFwdParamsWrapper &set_SoftmaxOffset(NVTETensor val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsSoftmaxOffset, val); } + FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsS, val); + } + FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsO, val); + } + FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack *val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAuxCtxTensors, val); + } FusedAttnFwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQ, val); } @@ -1233,30 +1247,42 @@ class FusedAttnFwdParamsWrapper { FusedAttnFwdParamsWrapper &set_rng_state(NVTETensor val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsRngState, val); } - FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsS, val); - } - FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsO, val); + FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenQ, val); } - FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack *val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsAuxCtxTensors, val); + FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenKV, val); } FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsIsTraining, static_cast(val)); } + FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); + } FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsCudaGraph, static_cast(val)); } - FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); + FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); } - FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); + FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsDropout, val); + } + FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQKVLayout, val); + } + FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsOFormat, val); + } + FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQKVScaleInvFormat, val); } FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsBiasType, val); } + FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); + } FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsSoftmaxType, val); } @@ -1269,27 +1295,6 @@ class FusedAttnFwdParamsWrapper { FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsBottomRightDiagonal, static_cast(val)); } - FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsDropout, val); - } - FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); - } - FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsQKVLayout, val); - } - FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsOFormat, val); - } - FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsQKVScaleInvFormat, val); - } - FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenQ, val); - } - FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenKV, val); - } FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { return set_attr(kNVTEFusedAttnFwdParamsWorkspace, val); } @@ -1310,11 +1315,11 @@ class FusedAttnFwdParamsWrapper { }; /*! \class FusedAttnBwdParamsWrapper - * \brief C++ helper for constructing an ``NVTEFusedAttnBwdParams``. + * \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``. + * 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: @@ -1398,36 +1403,18 @@ class FusedAttnBwdParamsWrapper { FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, val); } - FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsCudaGraph, static_cast(val)); - } - FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDeterministic, static_cast(val)); - } - FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); - } - FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); - } - FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsSoftmaxType, val); - } - FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsWindowSizeLeft, val); + FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenQ, val); } - FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsWindowSizeRight, val); + FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenKV, val); } - FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsBottomRightDiagonal, static_cast(val)); + FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); } FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { return set_attr(kNVTEFusedAttnBwdParamsDropout, val); } - FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); - } FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { return set_attr(kNVTEFusedAttnBwdParamsQKVLayout, val); } @@ -1446,11 +1433,29 @@ class FusedAttnBwdParamsWrapper { FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { return set_attr(kNVTEFusedAttnBwdParamsDOScaleInvFormat, val); } - FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenQ, val); + FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); } - FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenKV, val); + FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); + } + FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsSoftmaxType, val); + } + FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWindowSizeLeft, val); + } + FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWindowSizeRight, val); + } + FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsBottomRightDiagonal, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDeterministic, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCudaGraph, static_cast(val)); } FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { return set_attr(kNVTEFusedAttnBwdParamsWorkspace, val); diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index f8cd1308cb..a312d39b42 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -18,10 +18,14 @@ import transformer_engine_jax from transformer_engine_jax import ( + DType, JAXX_Scaling_Mode, + NVTE_Bias_Type, NVTE_Fused_Attn_Backend, + NVTE_Mask_Type, NVTE_QKV_Format, NVTE_QKV_Layout, + NVTE_Softmax_Type, ) from transformer_engine.jax.attention import ( AttnBiasType, @@ -138,6 +142,63 @@ class _FusedAttnConfig: ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. +@dataclass +class FusedAttnParams: + """ + Attention parameters used to select the fused attention backend. + + Fields are declared in the order of the ``FusedAttnConfig`` struct in + ``common/fused_attn/config_and_params.h``, which is the order the C++ binding reads them in + and the order it fills the config with. Fields JAX does not use, namely the paged-KV + dimensions and the ragged token counts, are omitted and keep their ``FusedAttnConfig`` + defaults. + """ + + # basic attention settings + is_training: bool = True + deterministic: bool = False + cuda_graph: bool = False + return_max_logit: bool = False + attn_mask_type: NVTE_Mask_Type = NVTE_Mask_Type.NVTE_NO_MASK + bias_type: NVTE_Bias_Type = NVTE_Bias_Type.NVTE_NO_BIAS + window_size_left: int = -1 + window_size_right: int = -1 + bottom_right_diagonal: bool = True + softmax_type: NVTE_Softmax_Type = NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX + scaling_mode: JAXX_Scaling_Mode = JAXX_Scaling_Mode.NO_SCALING + dropout: float = 0.0 + attn_scale: float = 1.0 + + # tensor types + qkv_dtype: DType = DType.kBFloat16 + o_dtype: DType = DType.kBFloat16 + do_dtype: DType = DType.kBFloat16 + dqkv_dtype: DType = DType.kBFloat16 + + # tensor layouts + qkv_layout: NVTE_QKV_Layout = NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + o_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + do_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + dqkv_layout: NVTE_QKV_Layout = NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + qkv_scale_inv_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + do_scale_inv_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + + # 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 + + # bias dimensions + bias_batch_size: int = 0 + bias_num_heads: int = 0 + bias_seqlen_q: int = 0 + bias_seqlen_kv: int = 0 + + @dataclass(frozen=True) class FusedAttnHelper: """ @@ -183,6 +244,9 @@ def get_fused_attn_backend(self): resolved config and the reason fused attention was rejected. """ q_type = jax_dtype_to_te_dtype(self.q_dtype) + kv_type = jax_dtype_to_te_dtype(self.kv_dtype) + if q_type != kv_type: + raise ValueError("Q and KV must have the same data type.") 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 or 0 @@ -190,39 +254,34 @@ def get_fused_attn_backend(self): 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, - q_type, - jax_dtype_to_te_dtype(self.kv_dtype), - q_type, - q_type, - q_type, - JAXX_Scaling_Mode.NO_SCALING, - self.qkv_layout.value, - NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET, - NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, - self.attn_bias_type.value, - self.attn_mask_type.value, - self.softmax_type.value, - self.attn_scale, - self.dropout_probability, - self.q_num_heads, - self.kv_num_heads, - self.q_max_seqlen, - self.kv_max_seqlen, - self.head_dim_qk, - self.head_dim_v, - self.window_size[0], - self.window_size[1], - self.bottom_right_diagonal, - not self.is_non_deterministic_allowed(), - bias_batch, - bias_heads, - bias_seqlen_q, - bias_seqlen_kv, + FusedAttnParams( + is_training=self.is_training, + deterministic=not self.is_non_deterministic_allowed(), + attn_mask_type=self.attn_mask_type.value, + bias_type=self.attn_bias_type.value, + window_size_left=self.window_size[0], + window_size_right=self.window_size[1], + bottom_right_diagonal=self.bottom_right_diagonal, + softmax_type=self.softmax_type.value, + dropout=self.dropout_probability, + attn_scale=self.attn_scale, + qkv_dtype=q_type, + o_dtype=q_type, + do_dtype=q_type, + dqkv_dtype=q_type, + qkv_layout=self.qkv_layout.value, + batch_size=self.batch_size, + num_attn_heads=self.q_num_heads, + num_gqa_groups=self.kv_num_heads, + head_dim_qk=self.head_dim_qk, + head_dim_v=self.head_dim_v, + max_seqlen_q=self.q_max_seqlen, + max_seqlen_kv=self.kv_max_seqlen, + bias_batch_size=bias_batch, + bias_num_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, + ) ) AttentionLogging.setup_logging() diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 905556d46a..b7ba1c6af5 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -152,17 +152,11 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); +// Select the fused attention backend for the configuration carried by a FusedAttnParams object +// (see jax/cpp_extensions/attention.py). Returns the backend and, when no backend supports the +// configuration, a diagnostic message explaining why it was rejected. std::tuple GetFusedAttnBackend( - bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, - DType do_dtype, DType dqkv_dtype, JAXX_Scaling_Mode scaling_mode, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - float 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, size_t bias_batch, size_t bias_heads, size_t bias_seqlen_q, - size_t bias_seqlen_kv); + const pybind11::object ¶ms); 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 30ff61b013..1c14f3d474 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -24,59 +24,48 @@ namespace transformer_engine { 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, JAXX_Scaling_Mode scaling_mode, NVTE_QKV_Layout qkv_layout, +static std::tuple GetFusedAttnBackendImpl( + bool is_training, bool deterministic, bool cuda_graph, bool return_max_logit, + NVTE_Mask_Type mask_type, NVTE_Bias_Type bias_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, NVTE_Softmax_Type softmax_type, + JAXX_Scaling_Mode scaling_mode, float dropout_probability, float attn_scale, DType q_dtype, + DType o_dtype, DType do_dtype, DType dqkv_dtype, 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, 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); - } - if (do_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { - do_format = o_format; - } - if (dqkv_layout == NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET) { - dqkv_layout = qkv_layout; - } - NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); - + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, size_t batch_size, + size_t q_attn_heads, size_t kv_attn_heads, size_t qk_head_dim, size_t v_head_dim, + size_t q_max_seqlen, size_t kv_max_seqlen, size_t bias_batch, size_t bias_heads, + size_t bias_seqlen_q, size_t bias_seqlen_kv) { 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_cuda_graph(cuda_graph) + .set_return_max_logit(return_max_logit) .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_bias_type(bias_type) .set_window_size_left(window_size_left) .set_window_size_right(window_size_right) .set_bottom_right_diagonal(bottom_right_diagonal) + .set_softmax_type(softmax_type) + .set_scaling_mode(get_nvte_scaling_mode(scaling_mode)) + .set_dropout(dropout_probability) + .set_attn_scale(attn_scale) .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_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_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) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) .set_bias_batch_size(bias_batch) .set_bias_num_heads(bias_heads) .set_bias_seqlen_q(bias_seqlen_q) @@ -87,6 +76,46 @@ std::tuple GetFusedAttnBackend( return {backend, message != nullptr ? std::string(message) : std::string()}; } +std::tuple GetFusedAttnBackend( + const pybind11::object ¶ms) { + const auto qkv_layout = params.attr("qkv_layout").cast(); + auto o_format = params.attr("o_format").cast(); + auto do_format = params.attr("do_format").cast(); + auto dqkv_layout = params.attr("dqkv_layout").cast(); + 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; + } + + return GetFusedAttnBackendImpl( + params.attr("is_training").cast(), params.attr("deterministic").cast(), + params.attr("cuda_graph").cast(), params.attr("return_max_logit").cast(), + params.attr("attn_mask_type").cast(), + params.attr("bias_type").cast(), + params.attr("window_size_left").cast(), + params.attr("window_size_right").cast(), + params.attr("bottom_right_diagonal").cast(), + params.attr("softmax_type").cast(), + params.attr("scaling_mode").cast(), + params.attr("dropout").cast(), params.attr("attn_scale").cast(), + params.attr("qkv_dtype").cast(), params.attr("o_dtype").cast(), + params.attr("do_dtype").cast(), params.attr("dqkv_dtype").cast(), qkv_layout, + o_format, do_format, dqkv_layout, + params.attr("qkv_scale_inv_format").cast(), + params.attr("do_scale_inv_format").cast(), + params.attr("batch_size").cast(), params.attr("num_attn_heads").cast(), + params.attr("num_gqa_groups").cast(), params.attr("head_dim_qk").cast(), + params.attr("head_dim_v").cast(), params.attr("max_seqlen_q").cast(), + params.attr("max_seqlen_kv").cast(), params.attr("bias_batch_size").cast(), + params.attr("bias_num_heads").cast(), params.attr("bias_seqlen_q").cast(), + params.attr("bias_seqlen_kv").cast()); +} + /* NOTE: PrepareFusedAttnForwardAuxTensors unifies the auxiliary tensor pack logic from the fused attention forward kernels in: @@ -349,14 +378,14 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ 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, 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); + auto [backend, _fwd_msg] = GetFusedAttnBackendImpl( + is_training, deterministic, false, false, mask_type, bias_type, window_size_left, + window_size_right, bottom_right_diagonal, softmax_type, JAXX_Scaling_Mode::NO_SCALING, + dropout_probability, scaling_factor, dtype, dtype, dtype, dtype, qkv_layout, + nvte_get_q_format(qkv_layout), nvte_get_q_format(qkv_layout), qkv_layout, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + input_batch, attn_heads, num_gqa_groups, qk_head_dim, v_head_dim, q_max_seqlen, kv_max_seqlen, + 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) */ @@ -680,14 +709,14 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ 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, 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); + auto [backend, _bwd_msg] = GetFusedAttnBackendImpl( + is_training, deterministic, false, false, mask_type, bias_type, window_size_left, + window_size_right, bottom_right_diagonal, softmax_type, JAXX_Scaling_Mode::NO_SCALING, + dropout_probability, scaling_factor, dtype, dtype, dtype, dtype, qkv_layout, + nvte_get_q_format(qkv_layout), nvte_get_q_format(qkv_layout), qkv_layout, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + input_batch, attn_heads, num_gqa_groups, qk_head_dim, v_head_dim, q_max_seqlen, kv_max_seqlen, + 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/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 78f4c9493c..bc31a74d4f 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -134,7 +134,8 @@ pybind11::dict Registrations() { PYBIND11_MODULE(transformer_engine_jax, m) { m.def("registrations", &Registrations); - m.def("get_fused_attn_backend", &GetFusedAttnBackend); + m.def("get_fused_attn_backend", &GetFusedAttnBackend, "Get Fused Attention backend", + pybind11::arg("fused_attn_params")); m.def("get_cuda_version", &GetCudaRuntimeVersion); m.def("get_cudnn_version", &GetCudnnRuntimeVersion); m.def("get_cudnn_frontend_version", &GetCudnnFrontendVersion); diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 9bafc9f500..b0d24ac1bc 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -2023,14 +2023,15 @@ def backward(ctx, d_out, *_args): class FusedAttention(torch.nn.Module): - """Dot product attention using cuDNN attention: + """Dot product attention using `cuDNN attention `_: FusedAttnBackend["F16_arbitrary_seqlen"] cuDNN attention for FP16/BF16 with any sequence length. FusedAttnBackend["FP8"] - 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. + cuDNN attention for FP8 with any sequence length. It supports the following recipes, where + "Inputs", "Intermediates" and "Outputs" are in the format of "tensor: quantizer". The recipes + are implemented in transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_fwd and + transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_bwd. Direction Inputs Intermediates Outputs DelayedScaling (DS) forward Q/K/V: DS S: DS O: DS @@ -2041,6 +2042,10 @@ class FusedAttention(torch.nn.Module): 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 + + For MXFP8, "row" and "col" are the quantization directions, which align with the contraction axes + of the matmuls that consume the tensor. For more details, please refer to + `How Scales Are Applied in MXFP8 Attention `_. """ def __init__( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index d3143281df..a1774c9df2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1665,7 +1665,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # symbolic (assume_constant_result requires concrete values); ints/floats made # dynamic by automatic dynamic currently graph break here. fused_attention_backend, reject_message = _get_fused_attn_backend(**fused_attn_kwargs) - if fused_attention_backend == FusedAttnBackend.No_Backend.value: + if fused_attention_backend == FusedAttnBackend.No_Backend: logger.debug( "Disabling FusedAttention: %s%s", reject_message, @@ -1682,7 +1682,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if ( use_fused_attention and has_score_mod - and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value + and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " @@ -1733,7 +1733,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend.FP8.value + fused_attention_backend == FusedAttnBackend.FP8 and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1744,7 +1744,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value + fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen and is_training and ( device_compute_capability < (9, 0) diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index c7272ee0a2..04f005522a 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -140,12 +140,13 @@ def cast( # Fail fast at import time if a new enumerator is added on the C++ side -# without being mirrored above. -assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set(NVTE_Fused_Attn_Backend.__members__), ( - "FusedAttnBackend in python is out of sync with" - " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." - " Please make sure TE C++ and python are in sync." -) +# without being mirrored above. Not an assert, so that the check survives -O. +if {f"NVTE_{m.name}" for m in FusedAttnBackend} != set(NVTE_Fused_Attn_Backend.__members__): + raise RuntimeError( + "FusedAttnBackend in python is out of sync with" + " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." + " Please make sure TE C++ and python are in sync." + ) BACKEND_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 5bfbc79d53..767c93140f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -54,6 +54,7 @@ std::tuple get_fused_attn_backend(const py .set_softmax_type(p.attr("softmax_type").cast()) .set_scaling_mode(p.attr("scaling_mode").cast()) .set_dropout(p.attr("dropout").cast()) + .set_attn_scale(p.attr("attn_scale").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())) @@ -64,7 +65,6 @@ std::tuple get_fused_attn_backend(const py .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()) From 76092d14e63baff61f186adf16deab3e892627ed Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:48:11 -0700 Subject: [PATCH 87/88] WIP: trim code/comments v2 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../fused_attn_f16_arbitrary_seqlen.cu | 34 +-- .../common/fused_attn/fused_attn_fp8.cu | 65 +----- .../common/fused_attn/graph_cache.h | 121 ++-------- .../common/fused_attn/graph_cache_debug.h | 212 ++---------------- 4 files changed, 39 insertions(+), 393 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 33f0ef52a1..175c1c327c 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 @@ -48,16 +48,6 @@ using F16FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// Constructs the forward graph for one cache key, and only constructs it: whether cuDNN will run -// it is settled by the caller, in cache_graph(), which is also where the plan build eventually -// happens. Hence no cuDNN handle here -- describing a graph needs none, and every call that does -// need one now sits on the other side of that boundary. -// -// Everything the graph's shape and topology depends on comes from `cfg`, so the build has one -// source of truth and cannot drift from the caller that will bind pointers to it. The two -// dimensions that differ between the passes -- the batch size, and the width ragged offsets are -// written in -- are read from the config's forward halves, the same ones the code binding pointers -// to this graph reads. static F16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg) { const int64_t b = static_cast(cfg.graph_batch_size_fwd); const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); @@ -371,21 +361,14 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - // Read from the same halves of the config the graph was built from, so that the dimensions below - // and the ones the graph was built at cannot be decided differently. Asserted derived here - // because these are the first derived fields this path reads, ahead of the get_graph() that - // asserts it for the build. cfg.check_derived(); const int64_t b = static_cast(cfg.graph_batch_size_fwd); const DType ragged_offset_type = cfg.ragged_offset_type_fwd; - // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by whatever the - // bucketing above did to `b`. const int64_t actual_b = static_cast(cfg.batch_size); const bool use_ragged_stats = cfg.uses_ragged_stats; const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; const bool return_max_logit = cfg.return_max_logit; - // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; @@ -575,8 +558,6 @@ using F16BwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The backward counterpart of create_graph_f16_fwd; see there for why it constructs the graph and -// nothing else, and why the two direction-dependent dimensions come from the config in pairs. static F16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg) { const int64_t b = static_cast(cfg.graph_batch_size_bwd); const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); @@ -849,18 +830,12 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cudnnHandle_t handle) { using namespace transformer_engine; - // Read from the same halves of the config the graph was built from, so that the dimensions below - // and the ones the graph was built at cannot be decided differently. Asserted derived here - // because these are the first derived fields this path reads, ahead of the get_graph() that - // asserts it for the build. cfg.check_derived(); const int64_t b = static_cast(cfg.graph_batch_size_bwd); const DType ragged_offset_type = cfg.ragged_offset_type_bwd; - // The true batch size, which the cu_seqlens buffers are sized [actual_b + 1] by. const int64_t actual_b = static_cast(cfg.batch_size); const bool use_ragged_stats = cfg.uses_ragged_stats; - // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; const bool is_bias = cfg.is_bias; const bool is_padding = cfg.is_padding; @@ -1044,8 +1019,6 @@ void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *i size_t i = 0; if (Aux_CTX_Tensors->size == 0) { - // These have to match the shape the forward graph declares for Stats and Max, which is why - // both read the same derived field rather than recomputing the condition. const bool use_ragged_stats = cfg.uses_ragged_stats; Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); @@ -1209,12 +1182,7 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i } } -// The one entry point into this translation unit's support probes; see fused_attn::support_verdict, -// which is all it does. It exists because create_graph_f16_* is file-local, so this is the only -// place that can name it, and because the selector calls it from another translation unit. -// -// Turning `pass` into the template argument is the whole of the body: the direction has to be a -// compile-time constant to pick a builder, and this is where the two meet. +// Check whether cuDNN can support a given config, per forward/backward pass. std::string support_verdict_f16(const FusedAttnConfig &cfg, Pass pass, cudnnHandle_t handle) { if (pass == Pass::Fwd) { return fused_attn::support_verdict(cfg, handle); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index eeec53207e..8d05a53ecf 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -20,7 +20,6 @@ namespace fused_attn { using namespace transformer_engine; namespace fe = cudnn_frontend; -// fused attention FWD FP8 with FE 1.0+ using Fp8FwdGraphAndTensors = std::tuple, std::shared_ptr, // Q @@ -44,43 +43,6 @@ using Fp8FwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// The three recipes these graphs are written for, read from cfg at each of the four sites that -// build or bind one: -// -// is_mxfp8 = scaling_mode is MXFP8 -// is_tensor_scaling = scaling_mode is DELAYED_TENSOR_SCALING -// is_delayed_scaling_fwd / _bwd = is_tensor_scaling && O / dQKV is FP8 -// is_current_scaling_fwd / _bwd = is_tensor_scaling && O / dQKV is F16 -// is_mxfp8_fwd / _bwd = is_mxfp8 && O / dQKV is F16 -// -// so at most one of the three holds for a pass, no combination of the booleans being able to say -// two things at once, and is_tensor_scaling is the delayed/current pair together -- which is what -// most of the sites below want, since a per-tensor scale is a per-tensor scale whichever recipe -// put it there. -// -// Each flag pairs a recipe with an output dtype it can write, so an output none of them can write -// leaves all three false rather than defaulting to one. That is the form the check takes in -// nvte_get_fused_attn_backend_v2, which refuses such a config before any graph here is built -- -// which is in turn why the sites below can treat the three as a partition. -// -// The split is per pass because a forward graph writes O and a backward one dQKV, and it is drawn -// on the output dtype because NVTEScalingMode has no current-scaling enumerator: both -// tensor-scaling recipes arrive as DELAYED_TENSOR_SCALING, and what separates them is whether the -// scale is known before the graph is built. See nvte_get_fused_attn_backend_v2 for the rest of -// what these graphs cannot represent, and config_and_params.h for the fields. -// -// Unlike the F16 path there is no bucketing to do, because FP8 has no ragged/THD support: the -// graph's shapes are exactly the config's. - -// Constructs the forward FP8 graph for one cache key, and only constructs it: whether cuDNN will -// run it is settled by the caller, in cache_graph(), which is also where the plan build eventually -// happens. Hence no cuDNN handle here -- describing a graph needs none, and every call that does -// need one now sits on the other side of that boundary. -// -// Everything the graph's shape and topology depends on comes from `cfg`, so the build has one -// source of truth and cannot drift from the caller that will bind pointers to it -- including the -// forward half of the pass-indexed fields, read here the same way the code binding pointers to -// this graph reads it. static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t qkv_tensor_type = @@ -365,18 +327,13 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - // Asserted derived here because the reads below are the first derived fields this path touches, - // ahead of the get_graph() that asserts it for the build. cfg.check_derived(); - // Read from the same fields the graph was built from, so that the tensors bound below and the - // ones the graph was built with cannot be decided differently. const bool is_tensor_scaling = cfg.is_tensor_scaling; const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; const bool use_cu_seqlens_directly = cfg.fp8_uses_cu_seqlens_directly; const int64_t b = static_cast(cfg.batch_size); - // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; const bool is_padding = cfg.is_padding; const bool is_dropout = cfg.is_dropout; @@ -467,7 +424,6 @@ void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* de } } -// fused attention BWD FP8 with FE 1.0+ using Fp8BwdGraphAndTensors = std::tuple, std::shared_ptr, // Q @@ -512,14 +468,6 @@ using Fp8BwdGraphAndTensors = std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset -// Builds the backward FP8 graph for one cache key, up to check_support() but not -// graph.build_plans(); see CacheEntry for why the plan build is left to whoever runs the graph. -// -// Everything the graph's shape and topology depends on is re-derived from `cfg` here, so the -// build has one source of truth for them. Unlike the F16 path, FP8 has no ragged/THD support, -// so the shapes are exactly the config's and need no bucketing from the caller. -// The backward counterpart of create_graph_fp8_fwd; see there for why it constructs the graph -// and nothing else. static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); const cudnn_frontend::DataType_t qkv_tensor_type = @@ -558,8 +506,6 @@ static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { const bool is_tensor_scaling = cfg.is_tensor_scaling; const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; const bool is_current_scaling = cfg.is_current_scaling_bwd; - // Whether O arrived in F16 rather than FP8, which decides whether this graph has to descale it on - // the way in. Read off O, unlike the recipe above, because O is what the forward pass stored. const bool is_O_in_F16 = !cfg.is_o_in_fp8; auto mha_graph = std::make_shared(); @@ -949,12 +895,8 @@ void fused_attn_fp8_bwd_impl( cudnnHandle_t handle) { using namespace transformer_engine; - // Asserted derived here because the reads below are the first derived fields this path touches, - // ahead of the get_graph() that asserts it for the build. cfg.check_derived(); - // Read from the same fields the graph was built from, so that the tensors bound below and the - // ones the graph was built with cannot be decided differently. const bool is_mxfp8 = cfg.is_mxfp8; const bool is_tensor_scaling = cfg.is_tensor_scaling; const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; @@ -962,7 +904,6 @@ void fused_attn_fp8_bwd_impl( const bool is_O_in_F16 = !cfg.is_o_in_fp8; const int64_t b = static_cast(cfg.batch_size); - // Not const: bound into the variant pack by address as a pass-by-value graph input. float scaling_factor = cfg.attn_scale; const bool is_padding = cfg.is_padding; const bool is_dropout = cfg.is_dropout; @@ -1305,11 +1246,7 @@ void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const } } -// The FP8 counterpart of support_verdict_f16; see there for why the direction arrives at runtime. -// -// Only cuDNN's rules reach this. TE's own -- bias, ALiBi and the recipes these graphs are not -// written for -- are stated in nvte_get_fused_attn_backend_v2 and answered before it ever gets -// here, which is why no build on this path throws for a configuration it cannot serve. +// Check whether cuDNN can support a given config, per forward/backward pass. std::string support_verdict_fp8(const FusedAttnConfig& cfg, Pass pass, cudnnHandle_t handle) { if (pass == Pass::Fwd) { return fused_attn::support_verdict(cfg, handle); diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h index 21c5098ab5..a694e7590b 100644 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -47,13 +47,6 @@ namespace transformer_engine { namespace fused_attn { -// A graph in the cache, plus the tensor attributes needed to bind runtime pointers to it. -// -// Entries stop at check_support(), which is all it takes to decide support. graph.build_plans() -- -// the kernel compilation, and the most expensive frontend call -- is left to the execution path, -// since a support query never runs the graph. build_plans_once guards that completion, which must -// happen exactly once per entry: the entry is shared across threads and graph.build_plans() mutates -// it in place. A build that throws leaves the flag unset, so a later call retries. template struct CacheEntry { explicit CacheEntry(GraphAndTensors graph_and_tensors) @@ -63,37 +56,12 @@ struct CacheEntry { std::once_flag build_plans_once; }; -// One build site's cache, process-wide rather than per-thread so a graph is reused across threads -// instead of rebuilt by each: cuDNN >= 9.0 allows concurrent execution of a shared plan, and the -// frontend's execute() builds its variant pack in a local rather than in the graph. A graph and its -// plans are compiled artifacts bound to the device they were finalized against, with nothing in -// them belonging to the building thread, which is why the key stamps device_id and nothing -// thread-shaped (see make_cache_key). -// -// The mutex is declared first so that it is destroyed last, leaving the map destroyed while its -// guard is still valid. -// -// The map is unbounded: a probe-only entry holds no compiled kernels, none hold a workspace, and a -// model reuses a handful of configurations. A workload that does sweep shapes holds every graph for -// the life of the process, which `miss` climbing without settling is the way to see. template struct GraphCache { - std::mutex mutex; // guards everything below + std::mutex mutex; std::map>> entries; }; -// Every cuDNN frontend call except graph.execute() runs holding this. The frontend serializes none -// of them for us, so two threads building unrelated keys is a data race, not the harmless duplicate -// work the map's view suggests. Not a theoretical exposure: a PyTorch step runs the forward and the -// backward's support probe on the main thread and the backward itself on the autograd thread. -// -// One lock for the process rather than one per cache, since what is unsafe is the frontend rather -// than any single graph. Kept separate from GraphCache::mutex, which guards only the map, so a hit -// never waits behind somebody else's kernel compilation. -// -// Lock ordering, which a later edit has to preserve: always taken before GraphCache::mutex, never -// after, and never held on entry to build_plans() -- holding it while waiting on that once_flag -// would deadlock against the thread holding the flag. inline std::mutex &frontend_build_mutex() { static std::mutex mutex; return mutex; @@ -103,16 +71,13 @@ inline std::mutex &frontend_build_mutex() { // `backend` and `pass` only name the build site the stage timers attribute the calls to. // // Reports by throwing, carrying cuDNN's message alone: that message is what support_verdict() -// returns as the reason a backend was refused, so a bool would discard the one thing a probe exists -// to produce, and NVTE_ERROR would dress a plain refusal as an internal failure. +// returns as the reason a backend was refused. inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, cudnnHandle_t handle) { auto run = [&](graph_cache_debug::BuildStage stage, const char *call_name, auto &&call) { const cudnn_frontend::error_t error = graph_cache_debug::record_time(backend, pass, stage, [&] { return call(); }); if (error.is_good()) return; - // cuDNN normally explains itself; fall back to the call's name so that a refusal can never - // arrive as an empty string, which support_verdict() would read as an endorsement. throw std::runtime_error(error.err_msg.empty() ? std::string(call_name) + " failed." : error.err_msg); }; @@ -126,23 +91,10 @@ inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Gra [&] { return graph.check_support(); }); } -// The cached entry for `key`, building and inserting it via `build` if absent: -// +// Cache for the entry `key`; build it first if absent. Record the lookup result and return the entry. // hit -> record HIT, return the entry // miss -> take frontend_build_mutex(), look again (a thread that raced us has finished by now), // record MISS, build(), query_support(), insert -// -// `build` only constructs a graph; this is what puts it through query_support(), so the cache holds -// exactly the graphs cuDNN agreed to run, and a hit skips those calls. A refusal throws and stores -// nothing, so the next query for a refused key is refused again -- fine for a settled run, since -// the frameworks re-enter the selector only when the configuration changes. -// -// `key` must be make_cache_key(pass)'s output for the same `pass`, not a raw execution config: two -// configs differing only in a field no graph reads (attn_scale, say) have to reach the same entry. -// -// The second look keeps a lost race cheap -- the loser would otherwise hold the one build lock to -// produce a graph it drops on the next line -- and keeps exactly one HIT or MISS per call, so miss -// still equals create_graph. template std::shared_ptr> cache_graph(GraphCache &cache, const FusedAttnConfig &key, @@ -156,9 +108,6 @@ std::shared_ptr> cache_graph(GraphCachesecond : nullptr; }; - // Recorded after the lock is dropped, so writing a trace line cannot hold up threads querying - // other keys. The counters stay exact, but two lookups that raced can be recorded in the opposite - // order, so a level-2 trace is the set of lookups, not their sequence. if (std::shared_ptr> cached = find()) { graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Hit, key); return cached; @@ -169,61 +118,33 @@ std::shared_ptr> cache_graph(GraphCache>(build()); graph_cache_debug::record_create_graph(backend, pass); - // Every site's tensor tuple leads with its graph, the one element this needs; a tuple ordered - // otherwise fails to compile rather than quietly validating the wrong object. - // - // The two counters bracket this call deliberately: a graph cuDNN refuses throws here, having - // recorded its CREATE_GRAPH and never reaching CACHE_GRAPH, so the gap between those two columns - // is cuDNN's refusals alone. + query_support(backend, pass, *std::get<0>(entry->graph_and_tensors), handle); graph_cache_debug::record_cache_graph(backend, pass); - // The insert always takes: a thread racing this key would have had to hold the build lock to do - // it, and the look above already ruled that out. + std::lock_guard lock(cache.mutex); return cache.entries.insert({key, std::move(entry)}).first->second; } -// A backend's graph cache for one pass, and the only route to it. The cache is this instantiation's -// static local, so the callers naming one triple share one cache and each -// triple gets its own. -// -// `kCreateGraphFn` is a template parameter rather than a `CreateFn &&` argument on purpose: as a -// parameter it makes the creator part of the instantiation. Passed as an argument, each distinct -// lambda type would instantiate its own static cache, and the two call sites for a pass would -// quietly stop sharing entries. +// Each backend's graph cache per forward/backward pass, called by support_verdict() and the +// execution path. The cache is this instantiation's static local, so the callers naming one +// triple share one cache and each triple gets its own. template auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { static GraphCache cache; - // Asserted once here for both the key and the graph, which read the same derived fields. cfg.check_derived(); return cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, [&] { return kCreateGraphFn(cfg); }); } -// Whether cuDNN can run the graph this config asks for, in one direction: the empty string if it -// can, otherwise cuDNN's own account of why not. This is the whole of what support_verdict_f16 and -// support_verdict_fp8 do; they exist only to reach their own translation unit's graph builders, -// which is also where a runtime direction becomes the compile-time one this needs. -// -// The question is answered by building the graph, so there is no separate list of rules to keep in -// step with the builder, and the graph lands in the cache the execution path reads. -// -// Refusals and failures read alike, because CUDNN_BACKEND_API_FAILED -- raised for any non-success -// cudnnStatus_t -- cannot separate CUDNN_STATUS_NOT_SUPPORTED from CUDNN_STATUS_ALLOC_FAILED. -// Either way this backend cannot serve this call and the caller wants the message. -// -// The direction is named by the caller rather than read off the config, which has both -// check_for_*_support set and so cannot say which graph is being probed. +// Check whether cuDNN can support a given config, per forward/backward pass. +// Returns an empty string if can; otherwise, a diagnostic string for the reason. template std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - // Support is signalled by returning the empty string, so a refusal that arrived without a message - // of its own needs a label rather than reading as an endorsement. auto label = [] { return std::string("support_verdict<") + graph_cache_debug::backend_name(kBackend) + ", " + graph_cache_debug::pass_name(kPass) + ">"; @@ -240,25 +161,9 @@ std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { } } -// Runs graph.build_plans(), the plan build cache_graph() left undone, once per entry. Call only -// when the graph is about to be executed: a support query builds entries nothing ever runs, and -// kernel compilation is the most expensive frontend call. See CacheEntry. -// -// The once_flag settles which thread runs the build, not whether it may run alongside another one: -// graph.build_plans() is a frontend call like the rest, so it also needs frontend_build_mutex(). -// The lock is taken inside the call_once rather than around it, so an entry whose plans are already -// built stays on the flag's atomic fast path and never touches the process-wide lock. -// -// Splitting the build in two means the thread that finishes it is often not the one that started it -// -- a sizing call caches the graph, an autograd thread is first to run it. What makes that safe: -// graph.build_plans() needs no handle (the overload accepting one ignores it, working from the -// operation graph descriptor and the device properties, which is where the >= 1.25.0 frontend the -// build requires is load-bearing); the handle that built that descriptor outlives the build because -// TE never destroys cuDNN handles (cudnnExecutionPlanManager leaves HandleManager's Destroy -// parameter null, so handles leak by design); the descriptor was finalized for one device, which is -// why the key carries device_id; and graph.execute() uses the running thread's own handle, so a -// handle is never used by two threads at once, which is what cuDNN asks in return for letting them -// share a plan. +// Previous calls only create the graph, caches it if verified to be supported. This function +// compiles the kernels via graph.build_plans(). It is the most expensive frontend call, and +// done only once per cache entry. template void build_plans(Backend backend, Pass pass, CacheEntry &entry) { std::call_once(entry.build_plans_once, [&] { diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h index 8358e7e24d..7e5de19d58 100644 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -75,26 +75,11 @@ namespace graph_cache_debug { inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } -// The frontend calls that make up a build, in the order they run. `kCount` must stay last: it sizes -// the timing table, and detail::kStageNames is indexed by these values, so the two must stay in the -// same order. enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; - -// What a lookup found: an entry, or nothing. enum class LookupResult { Miss, Hit }; -// ============================================================================ -// Machinery, in the order an event travels through it: the gate, the site index, the counters, the -// line, the exit summary. Nothing outside this file names any of it. -// ============================================================================ namespace detail { -// ============================================================================ -// The gate: whether this process records anything, and how it names itself when it does. Every -// answer here is fixed for the life of the process and cached in an initialized-once static, so a -// disabled build pays one load and one branch per call site. -// ============================================================================ - // Verbosity parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG: 0=off, 1=events, 2=trace. inline int debug_level() { static const int lvl = [] { @@ -119,16 +104,12 @@ inline int launcher_rank() { return rank; } -// On at level >= 1, and only for the ranks the ":" suffix selects. Every rank writes to the -// same stderr and under data/tensor parallelism they run identical shapes, so emitting from all of -// them multiplies the volume by the world size to say the same thing; hence rank 0 only by default. -// Context parallelism is the case worth overriding for, its ranks running different subsets of the -// per-step regimes. +// On at level >= 1, and only for the ranks the ":" suffix selects. Rank 0 only by default. inline bool enabled() { static const bool on = [] { if (debug_level() < 1) return false; const int rank = launcher_rank(); - if (rank < 0) return true; // sole process, nothing to filter + if (rank < 0) return true; const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); const char *sep = (e != nullptr) ? std::strchr(e, ':') : nullptr; if (sep == nullptr) return rank == 0; @@ -147,14 +128,10 @@ inline bool enabled() { return on; } -// The gate on the per-lookup and per-execution trace lines. Tests enabled() rather than just the -// level, so the answer holds wherever it is asked: level 2 alone is true on a rank that emits -// nothing, which would make this read as "trace" on every rank in the job. +// The gate on the trace printouts. Tests enabled() rather than just the level. inline bool enabled_with_trace() { return enabled() && debug_level() >= 2; } -// Names the emitting rank, without which the ranks sharing one stderr would be indistinguishable. -// The tag carries its own trailing separator, so a run whose launcher exports no rank prints no -// empty column. +// Names the emitting rank. inline const std::string &rank_tag() { static const std::string *tag = [] { const int rank = launcher_rank(); @@ -164,25 +141,22 @@ inline const std::string &rank_tag() { return *tag; } -// Short thread IDs (0, 1, 2, ...) in assignment order, not identity: tid=0 is whichever thread -// touched this cache first, and the number means nothing outside this process. +// Short thread IDs (0, 1, 2, ...) in assignment order. tid=0 is whichever thread touched this cache first. 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 event counters and build timings. +// Registered at first use. On process exit, prints event counters and build timings. Only used by the summary handler. inline void register_summary_once(); -// Backend major, pass minor, so that the two passes of one backend are adjacent -- which is how the -// counter lines and the summary rows present them. +// Backend major, pass minor, so that the two passes of one backend are adjacent. constexpr size_t kSiteCount = 4; inline constexpr size_t site_index(Backend b, Pass p) { return (b == Backend::F16 ? 0u : 2u) + (p == Pass::Fwd ? 0u : 1u); } -// ============================================================================ // Cache event counters, one block per build site. Each name is both the event tag on the line that // records it and the column carrying its running total: // - create_graph: a graph constructed for a miss, counted before cuDNN is asked about it. @@ -196,29 +170,6 @@ inline constexpr size_t site_index(Backend b, Pass p) { // - hit: a lookup answered from the cache. Need not lead to an execution -- it can be a backend // availability check, or a workspace-sizing call, which has no tensors to run with. // - miss: a lookup the cache did not answer; triggers a build. -// -// Identities, holding by construction, so a violation is a bug in the cache or in the counting: -// - hit + miss = every lookup, one per call to cache_graph() (the function, not the column of the -// same name), which makes it the denominator for the rest. -// - miss >= create_graph >= cache_graph, each drop a build that threw. create_graph - -// cache_graph is what cuDNN refused, and is where a refusal shows up. miss - create_graph would -// be a backend refusing from inside its own build, which none does since TE's FP8 rules moved -// to nvte_get_fused_attn_backend_v2 to be answered with a reason instead, so that gap should -// read zero and the column stays as the thing that says so. -// - cache_graph >= build_plans, the gap being graphs a probe built that nothing has run. -// - execute > 0 implies build_plans > 0, every site building plans ahead of the workspace-sizing -// return, itself ahead of record_execute. So a sizing call pays build_plans and never execute. -// - Both build identities belong to the totals rows, not to one thread's: the thread that builds -// a graph need not compile its plans, and a PyTorch step splits exactly that way. -// - Per-thread rows sum column by column to "tid=all dev=all", and the per-backend rows of one -// pass to that pass's all-backends row. -// - Stage timing calls fall along validate >= build_operation_graph >= create_execution_plans >= -// check_support, each drop being builds that ended at the stage before, which localizes where -// cuDNN refuses rather than only how long refusing took. The build_plans timing row can exceed -// its column: the timer records while unwinding, the counter only on return. -// -// What the columns say about a workload is user-facing and lives in docs/envvars.rst. -// ============================================================================ struct EventCounters { std::atomic create_graph{0}; @@ -234,9 +185,7 @@ inline EventCounters &counters(Backend b, Pass p) { return table[site_index(b, p)]; } -// One counter block read out into plain values, so the summary can sum blocks for its per-backend -// and all-backends rows. Not read as one indivisible operation, which nothing here wants: the -// summary runs at exit, after the writing threads are done. +// One counter block read out into plain values, so the summary can sum blocks for its per-backend and all-backends rows. struct CounterSnapshot { uint64_t create_graph = 0; uint64_t cache_graph = 0; @@ -255,8 +204,7 @@ struct CounterSnapshot { return *this; } - // Whether this block saw nothing at all, which is what lets the summary leave out the rows for a - // backend the run never used rather than printing zeros for it. + // Whether this block saw nothing at all. bool empty() const { return (create_graph | cache_graph | build_plans | execute | hit | miss) == 0; } @@ -273,13 +221,7 @@ inline CounterSnapshot snapshot(const EventCounters &c) { return s; } -// Per-thread counters, so the summary can break every column down by thread and backend: in the -// single-process context-parallel case each device is driven by its own thread, and under PyTorch -// this separates the main thread from the autograd one. -// -// `device` is the device this thread last drove, restamped on every event. Event lines print the -// live current device instead, which is exact; this exists for the per-thread summary rows, written -// at exit by whichever thread is exiting. +// Per-thread counters, so the summary can break every column down by thread and backend. struct ThreadCounters { unsigned tid = 0; std::atomic device{-1}; @@ -306,8 +248,6 @@ inline ThreadCounters &thread_counters() { static thread_local ThreadCounters *tc = [] { auto *p = new ThreadCounters(); p->tid = thread_seq_id(); - // Stamped here as well as on every event, so a thread that only ever hits the cache -- never - // reaching print_counters() at level 1 -- still names a device rather than the -1 it began at. p->device.store(cuda::current_device(), std::memory_order_relaxed); { std::lock_guard lock(thread_registry_mutex()); @@ -322,17 +262,7 @@ inline EventCounters &thread_counters(Backend b, Pass p) { return thread_counters().sites[site_index(b, p)]; } -// ============================================================================ -// Turning a counter block into a line, and getting a line out. One formatter, shared by the event -// lines and the summary rows so the two cannot drift apart, and one writer, so everything here -// reaches stderr the same way. -// ============================================================================ - -// The one place diagnostics reach stderr, and the reason it exists: the first line this process -// writes carries a leading newline. Diagnostics share stderr with whatever the framework is -// printing, and a test runner's progress output has no trailing newline of its own, so without this -// the first line continues someone else's. Where the previous output did end cleanly the prefix -// reads as a blank line setting the diagnostics apart. +// The one place diagnostics reach stderr. inline void write_stderr(const std::string &text) { static std::atomic first_line{true}; if (first_line.exchange(false, std::memory_order_relaxed)) { @@ -344,15 +274,7 @@ inline void write_stderr(const std::string &text) { std::fflush(stderr); } -// Format one counter block -- one pass of one backend -- as one line. One pass rather than both -// because a line carrying the forward and backward columns together ran past 300 characters; the -// two passes are adjacent rows instead. -// -// `tid_field` and `dev_field` are whole columns, e.g. "tid=3" and "dev=0"; the totals rows pass -// "tid=all" and "dev=all". `label` is the build site, "f16 fwd", plus the event name on an event -// line, and arrives padded to the width its own kind of line uses -- 20 characters for an event -// line, 7 for a summary row -- deliberately not one width for both, which would put twelve blank -// columns on every summary row to align the scattered event lines. +// Format one counter block -- one pass of one backend -- as one line. inline std::string format_counter_line(const char *tid_field, const char *dev_field, const char *label, const CounterSnapshot &c) { char buf[512]; @@ -365,30 +287,20 @@ inline std::string format_counter_line(const char *tid_field, const char *dev_fi return std::string(buf); } -// One event line, from the thread the event happened on, carrying the running totals of the build -// site that raised it. The device is read live rather than remembered, so it is the device this -// event was actually issued against, and is recorded on the thread's block for the exit summary. +// One event line, from the thread the event happened on, carrying the running totals of the build site that raised it. inline void print_counters(Backend b, Pass p, const char *event) { const int device = cuda::current_device(); thread_counters().device.store(device, std::memory_order_relaxed); char label[32]; char tid_field[16]; char dev_field[16]; - // The event name is padded to the longest of them, so that the counters of one event line fall - // where the next one's do. std::snprintf(label, sizeof(label), "%s %s %-12s", backend_name(b), pass_name(p), event); std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); write_stderr(format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p)))); } -// The body every recorder shares: gate, register the exit summary, and add one to `column` in both -// the process-wide block and this thread's. Returns whether diagnostics are on at all, so a caller -// can skip building a line nobody will read. -// -// Both blocks or neither. Moving one and not the other would leave the per-thread rows failing to -// add up to the totals row, which the summary presents as an invariant, and the discrepancy would -// look like a threading bug in the cache rather than a miscount here. +// The body every recorder shares: gate, register the exit summary, and add one to `column` in both the process-wide block and this thread's. inline bool record_counter(Backend b, Pass p, std::atomic EventCounters::*column) { if (!enabled()) return false; register_summary_once(); @@ -397,8 +309,7 @@ inline bool record_counter(Backend b, Pass p, std::atomic EventCounter return true; } -// The column a lookup lands in, and the tag naming it. Both are switches with no default, so adding -// an outcome fails to compile here rather than being silently counted as a miss. +// The column a lookup lands in. inline std::atomic EventCounters::*lookup_column(LookupResult result) { switch (result) { case LookupResult::Hit: @@ -419,32 +330,16 @@ inline const char *lookup_name(LookupResult result) { return "MISS"; } -// ============================================================================ -// Graph build timings. Which stage dominates says what to do about a slow build: time in -// check_support and build_plans is heuristic selection and kernel compilation, largely intrinsic to -// the shape, while time in validate or build_operation_graph is graph-construction cost on our -// side. One duration per build could not make that distinction. -// -// Only sums are kept, so the summary reports a mean and nothing else -- and since a build happens -// once per distinct cache key, those calls span different shapes rather than repeating one. Read a -// stage mean as where build time goes in aggregate, not as any one build's cost. -// ============================================================================ - -// Indexed by BuildStage when the summary prints, so it must stay in that enum's order and carry one -// name per stage ahead of its kCount sentinel. inline constexpr const char *kStageNames[] = { "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; -// Totals for one (pass, stage) pair. Relaxed ordering is sufficient: these counters order nothing, -// and the only read happens once, after the threads that wrote them are done. +// Totals for one (pass, stage) pair. struct StageTiming { std::atomic calls{0}; std::atomic time_ns{0}; }; -// Bucketed by build site, an fp8 build and an f16 build being different work. Unlike the thread -// registry above, this needs no leak to outlive the exit handler that reads it: it holds nothing -// but atomics, so no destructor is registered for it at all. +// Bucketed by build site, fp8 vs f16. constexpr size_t kStageBuckets = kSiteCount * static_cast(BuildStage::kCount); inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { static std::array table{}; @@ -453,10 +348,7 @@ inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { return table[idx]; } -// Times one stage: clock read in the constructor, accumulated in the destructor. Recording on scope -// exit rather than at an explicit stop() keeps a failing stage measurable, since build_plans throws -// through NVTE_CHECK_CUDNN_FE and the destructor still runs while unwinding. `on` is latched at -// construction, so the destructor can never accumulate against an unset `start`. +// Times one stage: clock read in the constructor, accumulated in the destructor. struct ScopedBuildTimer { BuildStage stage; bool on; @@ -480,22 +372,14 @@ struct ScopedBuildTimer { } }; -// ============================================================================ -// Summary: on process exit, print cache event counters and graph build timings. Each piece below -// appends its rows to the block the handler is assembling, in the order they are printed. -// ============================================================================ - -// The two backends that keep a cache, in the order every part of the summary walks them. inline constexpr Backend kSummaryBackends[] = {Backend::F16, Backend::FP8}; -// Names one build site for a summary row. No padding: the site name is exactly the width of the -// column there, unlike the event lines, which pad it to keep their counters aligned. +// Names one build site for a summary row. inline std::string site_label(Backend b, Pass p) { return std::string(backend_name(b)) + " " + pass_name(p); } -// How many backends the run actually drove. Decides whether the across-backend rows are worth -// printing: with one backend they would repeat that backend's own rows verbatim. +// How many backends the run used. inline size_t active_backend_count() { size_t active = 0; for (const Backend b : kSummaryBackends) { @@ -506,8 +390,7 @@ inline size_t active_backend_count() { return active; } -// Per-thread breakdown, sorted by tid, one row per build site that thread drove. Sites it never -// reached are left out, for the reason an unused backend is: a row of zeros says nothing. +// Per-thread breakdown, sorted by tid, one row per build site that thread used. inline void append_thread_rows(std::string &block) { std::lock_guard lock(thread_registry_mutex()); std::vector blocks = thread_registry(); @@ -529,10 +412,7 @@ inline void append_thread_rows(std::string &block) { } } -// Totals, printed after the per-thread rows so they read as their sum: one row per build site, then -// one per pass across the backends when the run used more than one. Both come from the process-wide -// counters rather than by adding up the rows above, so the two agreeing is a check on the counting -// rather than an artifact of it. +// Totals, printed after the per-thread rows so they read as their sum: one row per build site, then one per pass across the backends when the run used more than one. inline void append_total_rows(std::string &block) { CounterSnapshot all_fwd; CounterSnapshot all_bwd; @@ -553,7 +433,7 @@ inline void append_total_rows(std::string &block) { } } -// Mean time per call for each stage of each build site, skipping stages nothing reached. +// Mean time per call for each stage of each build site, skipping stages that no build reached. inline void append_stage_rows(std::string &block) { for (const Backend b : kSummaryBackends) { for (const Pass p : {Pass::Fwd, Pass::Bwd}) { @@ -578,8 +458,7 @@ inline void register_summary_once() { static const bool registered = [] { std::atexit([] { if (!enabled()) return; - // Built in memory and emitted with one write, so that concurrently-exiting processes (one per - // rank under torchrun) stay grouped rather than interleaving. + // Built in memory and emitted with one write, so that concurrently-exiting processes stay grouped. const std::string marker = "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary "; std::string block = marker + "begin =====\n"; append_thread_rows(block); @@ -595,51 +474,27 @@ inline void register_summary_once() { } // namespace detail -// ============================================================================ // The recorders: everything a call site calls. Each takes the build site it is reporting for, adds // one to that site's column, and prints a line when the level asks for it. -// -// Every one of them is called after the event it names, never before, so that a column counts what -// happened rather than what was attempted. That is what gives the gaps between columns their -// meaning: an event that fails partway leaves the earlier column moved and the later one not. -// record_time is the exception, and only because timing cannot be done after the fact. -// ============================================================================ -// A graph constructed for a miss, whatever cuDNN goes on to make of it. Call as soon as -// construction returns and before check_support() is asked: construction is where a backend would -// refuse on its own rules, which is what makes miss - create_graph the builds that failed on this -// side of cuDNN. No backend does that now, so the gap is there to read as zero. inline void record_create_graph(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::create_graph)) { detail::print_counters(b, p, "CREATE_GRAPH"); } } -// A created graph that cleared check_support(). Call as soon as that verdict returns, ahead of the -// insert: a refused graph throws in between, leaving its CREATE_GRAPH unanswered, which is what -// makes create_graph - cache_graph cuDNN's refusals. inline void record_cache_graph(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::cache_graph)) { detail::print_counters(b, p, "CACHE_GRAPH"); } } -// The graph.build_plans() a cache_graph deferred, now completed. Call from inside the call_once -// that runs it, and after the call returns: it throws without setting the once_flag, leaving a -// later execution to retry, so counting on the way out keeps this a count of runnable graphs. inline void record_build_plans(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::build_plans)) { detail::print_counters(b, p, "BUILD_PLANS"); } } -// An execution cuDNN accepted. Call after graph.execute() returns, so a graph the surrounding -// setup never reached is not counted as having run -- the stream set and the cu_seqlens conversion -// kernels sit in between, and either can throw. Accepted is as far as this can go: execute() -// enqueues and returns, so a fault the device raises later still leaves the execution counted here. -// -// Unlike the recorders above, this fires on every execution rather than once per distinct key, so -// its line is held back to level 2 while its column keeps counting. inline void record_execute(Backend b, Pass p) { if (detail::record_counter(b, p, &detail::EventCounters::execute) && detail::enabled_with_trace()) { @@ -647,15 +502,7 @@ inline void record_execute(Backend b, Pass p) { } } -// `key` is the normalized cache key -- make_cache_key(pass)'s output, the exact value looked up -- -// not the execution config it came from. HIT/MISS is decided by comparing keys, so a trace of -// anything else could not explain its own outcome, and diffing two MISS lines here names exactly -// the fields responsible for the extra build. The cost is that overwritten fields no longer appear -// in their original form: attn_scale reads 1, ragged num_tokens read 0, max_seqlen and batch_size -// read their bucketed values. inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedAttnConfig &key) { - // The highest-volume line here, one per cache lookup; keep it out of the level-1 path and off the - // stderr lock unless tracing. if (!detail::record_counter(b, p, detail::lookup_column(result)) || !detail::enabled_with_trace()) { return; @@ -665,20 +512,9 @@ inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedA "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | ", detail::rank_tag().c_str(), detail::thread_seq_id(), key.device_id, backend_name(b), pass_name(p), detail::lookup_name(result)); - // The one line here not built from counters: which fields it names is - // FusedAttnConfig::to_string()'s to say, alongside the operator< that decides what a key - // compares on in the first place. detail::write_stderr(prefix + key.to_string() + "\n"); } -// Record how long `fn` takes as `stage` of the given build site. Unlike the recorders above this -// wraps the work rather than reporting on work already done, so the measured region is exactly the -// call passed in. Stage timings feed the summary only; they print no line of their own. -// -// Passes `fn`'s result back out so that a timed call reporting a value can be written as the -// initializer of that value. cuDNN's error_t is [[nodiscard]], and the alternative -- declaring the -// variable above the timing and assigning to it inside a void `fn` -- discards the assignment's own -// result, which the compiler counts as ignoring a nodiscard value. template inline decltype(auto) record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { detail::ScopedBuildTimer scoped(b, p, stage); From b1417ca4065b3c97ecf83c085f99083bec293aee Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:49:30 +0000 Subject: [PATCH 88/88] [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 | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 1c14f3d474..f24900410e 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -101,11 +101,10 @@ std::tuple GetFusedAttnBackend( params.attr("window_size_right").cast(), params.attr("bottom_right_diagonal").cast(), params.attr("softmax_type").cast(), - params.attr("scaling_mode").cast(), - params.attr("dropout").cast(), params.attr("attn_scale").cast(), - params.attr("qkv_dtype").cast(), params.attr("o_dtype").cast(), - params.attr("do_dtype").cast(), params.attr("dqkv_dtype").cast(), qkv_layout, - o_format, do_format, dqkv_layout, + params.attr("scaling_mode").cast(), params.attr("dropout").cast(), + params.attr("attn_scale").cast(), params.attr("qkv_dtype").cast(), + params.attr("o_dtype").cast(), params.attr("do_dtype").cast(), + params.attr("dqkv_dtype").cast(), qkv_layout, o_format, do_format, dqkv_layout, params.attr("qkv_scale_inv_format").cast(), params.attr("do_scale_inv_format").cast(), params.attr("batch_size").cast(), params.attr("num_attn_heads").cast(),