From a1ff90b34fdb5588cdee78d1df82065745f39a4c Mon Sep 17 00:00:00 2001 From: Xinya Zhang Date: Wed, 27 May 2026 15:23:47 +0000 Subject: [PATCH] [ROCm] Bump AOTriton to 0.12b (#184288) Notable new features: * **BREAKING** Varlen LSE tensor shape changes to (H, Total_seqlen) * Support head_dim != head_dim_v * Support `use_deterministic_algorithims` * Support seqused_k in test/test_varlen_attention.py * gfx1100 and gfx1151 promoted out of experimental * Partial FAv3 support on gfx950 Bug Fixes: * GQA kernel failed to read bias tensor with the right offset. Known Issues * gfx950's Triton kernel has problem handling hdim=16's fwd, in addition to hdim=48/80's bwd. * Disables gfx90a's CK SDPA support due to GPU Segfault. Pull Request resolved: https://github.com/pytorch/pytorch/pull/184288 Approved by: https://github.com/jeffdaily --- aten/src/ATen/Context.cpp | 2 +- .../native/transformers/cuda/attention.cu | 129 ++---- .../transformers/cuda/attention_backward.cu | 190 +++----- .../native/transformers/cuda/sdp_utils.cpp | 78 +++- .../transformers/hip/aotriton_adapter.h | 17 +- .../transformers/hip/aotriton_versions.h | 6 + .../hip/flash_attn/aot/mha_all_aot.hip | 417 +++++++----------- cmake/External/aotriton.cmake | 37 +- test/test_transformers.py | 8 - test/test_varlen_attention.py | 5 +- torch/nn/attention/varlen.py | 17 - 11 files changed, 356 insertions(+), 550 deletions(-) diff --git a/aten/src/ATen/Context.cpp b/aten/src/ATen/Context.cpp index c342590b58c42..d0c1d6eb88887 100644 --- a/aten/src/ATen/Context.cpp +++ b/aten/src/ATen/Context.cpp @@ -505,7 +505,7 @@ at::BlasBackend Context::blasPreferredBackend() { bool Context::ckSupported() { #ifdef USE_ROCM static const std::vector supported_archs = { - "gfx90a", "gfx942", "gfx950" + "gfx942", "gfx950" }; for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) { if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) { diff --git a/aten/src/ATen/native/transformers/cuda/attention.cu b/aten/src/ATen/native/transformers/cuda/attention.cu index 69ecf31df0586..93466a28064f4 100644 --- a/aten/src/ATen/native/transformers/cuda/attention.cu +++ b/aten/src/ATen/native/transformers/cuda/attention.cu @@ -1563,6 +1563,10 @@ std::tuple _efficient_ // compute_logsumexp is false constexpr int kAlignLSE = 1; res = at::empty({B, M, num_heads, Kv}, query.options()); + // TODO: Use Compact Varlen LSE + // The current memory allocation is strictly larger than necessary + // (total_q <= max_seqlen_q * B) + // The problem is total_q is not available here. at::Tensor softmax_lse; logsumexp = at::empty( { B, num_heads, compute_logsumexp ? max_seqlen_q : 0}, @@ -1591,8 +1595,6 @@ std::tuple _efficient_ atomic_counter = at::zeros({1}, query.options().dtype(at::kInt)); } - using aotriton::v2::flash::attn_fwd; - using aotriton::v2::flash::attn_fwd_compact_varlen; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::mk_philoxtensor; @@ -1608,92 +1610,47 @@ std::tuple _efficient_ auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr() : nullptr); auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); hipError_t err; // TODO: Error handling - if constexpr (AOTRITON_ALWAYS_V3_API) { // Better readability than nesting ifdef -#if AOTRITON_V3_API // if constexpr does not stop errors from undefined functions - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - using aotriton::v3::flash::WindowValue; - aotriton::v3::flash::attn_fwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.Sm_scale = softmax_scale; - params.L = compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2; - params.Out = mk_aotensor(output_t, "Out"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = dropout_p; - params.philox_seed_ptr = seed; - params.philox_offset1 = offset1; - params.philox_offset2 = offset2; - params.philox_seed_output = seed_output; - params.philox_offset_output = offset_output; - params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); - params.persistent_atomic_counter = persistent_counter; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { - params.window_left = WindowValue::TopLeftAligned; - params.window_right = WindowValue::TopLeftAligned; - } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { - params.window_left = WindowValue::BottomRightAligned; - params.window_right = WindowValue::BottomRightAligned; - } - if (bias.has_value()) { - params.B = mk_aotensor(bias.value(), "bias"); - } - if (seqstart_q.has_value()) { - params.varlen_type = VarlenType::CompactVarlen; - params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"); - } else { - params.varlen_type = VarlenType::None; - } - err = aotriton::v3::flash::attn_fwd(params, - aotriton::v3::flash::attn_fwd_params::kVersion, - stream); -#endif // AOTRITON_V3_API - } else if (seqstart_q.has_value()) { - // varlen aka nested tensor - err = attn_fwd_compact_varlen(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias"): empty_t4, - mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"), - mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - softmax_scale, - compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2, - mk_aotensor(output_t, "Out"), - dropout_p, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + using aotriton::v3::flash::WindowValue; + aotriton::v3::flash::attn_fwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.Sm_scale = softmax_scale; + params.L = compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2; + params.Out = mk_aotensor(output_t, "Out"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = dropout_p; + params.philox_seed_ptr = seed; + params.philox_offset1 = offset1; + params.philox_offset2 = offset2; + params.philox_seed_output = seed_output; + params.philox_offset_output = offset_output; + params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); + params.persistent_atomic_counter = persistent_counter; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { + params.window_left = WindowValue::TopLeftAligned; + params.window_right = WindowValue::TopLeftAligned; + } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { + params.window_left = WindowValue::BottomRightAligned; + params.window_right = WindowValue::BottomRightAligned; + } + if (bias.has_value()) { + params.B = mk_aotensor(bias.value(), "bias"); + } + if (seqstart_q.has_value()) { + params.varlen_type = VarlenType::CompactVarlen; + params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"); } else { - err = attn_fwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias"): empty_t4, - softmax_scale, - compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2, - mk_aotensor(output_t, "Out"), - dropout_p, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); + params.varlen_type = VarlenType::None; } + err = aotriton::v3::flash::attn_fwd(params, + aotriton::v3::flash::attn_fwd_params::kVersion, + stream); #else TORCH_CHECK(false, "Attempting to use AOTriton mem_eff_forward backend in a build that has not built AOTriton"); #endif diff --git a/aten/src/ATen/native/transformers/cuda/attention_backward.cu b/aten/src/ATen/native/transformers/cuda/attention_backward.cu index 183f99e975cda..b7e52617a0ec6 100644 --- a/aten/src/ATen/native/transformers/cuda/attention_backward.cu +++ b/aten/src/ATen/native/transformers/cuda/attention_backward.cu @@ -547,6 +547,17 @@ _efficient_attention_backward( "[AOTriton] Accelerated SDPA only supports MI200/MI300X/7900XTX/9070XT GPUs" " (gfx90a/gfx942/gfx1100/gfx1201)") } + bool deterministic{false}; + auto& ctx = at::globalContext(); + if (ctx.deterministicAlgorithms()) { + if (ctx.deterministicAlgorithmsWarnOnly()) { + TORCH_WARN_ONCE( + "Memory Efficient attention defaults to a non-deterministic algorithm. ", + "To explicitly enable determinism call torch.use_deterministic_algorithms(True, warn_only=False)."); + } else { + deterministic = true; + } + } const auto softmax_scale = sdp::calculate_scale(query, scale).expect_float(); bool is_causal; if (static_cast(sdp::CustomMaskType::NoCustomMask) == custom_mask_type) { @@ -569,139 +580,60 @@ _efficient_attention_backward( at::Tensor dout_t = grad_out.permute({0,2,1,3}); at::Tensor softmax_lse = logsumexp.view({B * nH, max_seqlen_q}); hipError_t err; - using aotriton::v2::flash::attn_bwd; - using aotriton::v2::flash::attn_bwd_fused; - using aotriton::v2::flash::attn_bwd_compact_varlen; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::cast_dtype; aotriton::TensorView<4> empty_t4(0, {0, 0, 0, 0}, {0, 0, 0, 0}, cast_dtype(query.dtype())); - if constexpr (AOTRITON_ALWAYS_V3_API) { // Better readability than nesting ifdef -#if AOTRITON_V3_API // if constexpr does not stop errors from undefined functions - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - using aotriton::v3::flash::WindowValue; - aotriton::v3::flash::attn_bwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.B = bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4; - params.Sm_scale = softmax_scale; - params.Out = mk_aotensor(out_t, "out"); - params.DO = mk_aotensor(dout_t, "dout"); - params.DK = mk_aotensor(dk_t, "dk"); - params.DV = mk_aotensor(dv_t, "dv"); - params.DQ = mk_aotensor(dq_t, "dq"); - params.DB = bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4; - params.L = mk_aotensor<2>(softmax_lse, "L"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = float(dropout_p); - params.philox_seed_ptr = mk_aoscalartensor(philox_seed); - params.philox_offset1 = mk_aoscalartensor(philox_offset); - params.philox_offset2 = 0; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { - params.window_left = WindowValue::TopLeftAligned; - params.window_right = WindowValue::TopLeftAligned; - } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { - params.window_left = WindowValue::BottomRightAligned; - params.window_right = WindowValue::BottomRightAligned; - } -#if AOTRITON_ALWAYS_V3_API - using sdp::aotriton_adapter::mklazy_empty_like; - using sdp::aotriton_adapter::mklazy_fp32zeros; - using sdp::aotriton_adapter::LazyTensorContext; - LazyTensorContext lazy_delta { .like_tensor = softmax_lse, .tensor_name = "delta" }; - LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; - params.D = mklazy_empty_like<2>(&lazy_delta); - params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); -#else - at::Tensor delta = at::empty_like(softmax_lse).contiguous(); - params.D = mk_aotensor<2>(delta, "delta"); -#endif - if (cu_seqlens_q.has_value()) { - params.varlen_type = VarlenType::CompactVarlen; - params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"); - } else { - params.varlen_type = VarlenType::None; - } - err = aotriton::v3::flash::attn_bwd(params, - aotriton::v3::flash::attn_bwd_params::kVersion, - stream); -#endif // AOTRITON_V3_API - } else if (cu_seqlens_q.has_value()) { - at::Tensor delta = at::empty_like(softmax_lse).contiguous(); - // varlen aka Nested tensor - err = attn_bwd_compact_varlen(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"), - mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4, - mk_aotensor<2>(softmax_lse, "L"), - mk_aotensor<2>(delta, "delta"), - float(dropout_p), - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } else { // cu_seqlens.has_value - auto d_head = Kv; - bool use_fused_bwd = d_head <= 192 && d_head * max_seqlen_q < 64 * 512; - if (use_fused_bwd) { - err = attn_bwd_fused(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4, - mk_aotensor<2>(softmax_lse, "L"), - float(dropout_p), - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } else { - at::Tensor delta = at::empty_like(softmax_lse).contiguous(); - err = attn_bwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4, - mk_aotensor<2>(softmax_lse, "L"), - mk_aotensor<2>(delta, "delta"), - float(dropout_p), - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } //used_fused_bwd - } // cuseqlen.has_value + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + using aotriton::v3::flash::WindowValue; + aotriton::v3::flash::attn_bwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.B = bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4; + params.Sm_scale = softmax_scale; + params.Out = mk_aotensor(out_t, "out"); + params.DO = mk_aotensor(dout_t, "dout"); + params.DK = mk_aotensor(dk_t, "dk"); + params.DV = mk_aotensor(dv_t, "dv"); + params.DQ = mk_aotensor(dq_t, "dq"); + params.DB = bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4; + params.L = mk_aotensor<2>(softmax_lse, "L"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = float(dropout_p); + params.philox_seed_ptr = mk_aoscalartensor(philox_seed); + params.philox_offset1 = mk_aoscalartensor(philox_offset); + params.philox_offset2 = 0; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { + params.window_left = WindowValue::TopLeftAligned; + params.window_right = WindowValue::TopLeftAligned; + } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { + params.window_left = WindowValue::BottomRightAligned; + params.window_right = WindowValue::BottomRightAligned; + } + using sdp::aotriton_adapter::mklazy_empty_like; + using sdp::aotriton_adapter::mklazy_fp32zeros; + using sdp::aotriton_adapter::LazyTensorContext; + LazyTensorContext lazy_delta { .like_tensor = softmax_lse, .tensor_name = "delta" }; + LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; + params.D = mklazy_empty_like<2>(&lazy_delta); + params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); + if (cu_seqlens_q.has_value()) { + params.varlen_type = VarlenType::CompactVarlen; + params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"); + } else { + params.varlen_type = VarlenType::None; + } + aotriton::v3::flash::attn_options opts; + opts.deterministic = deterministic; + err = aotriton::v3::flash::attn_bwd(params, + aotriton::v3::flash::attn_bwd_params::kVersion, + stream, + &opts); #else // DISABLE_AOTRITON TORCH_CHECK(false, "Attempting to use aotriton mem_eff_backward backend in a build that has not built AOTriton"); #endif diff --git a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp index 79b5df3f302bb..41ad5bacf1ed4 100644 --- a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp +++ b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp @@ -140,19 +140,11 @@ int64_t minimum_gemm_alignment(sdp_params const& params) { return matmul_alignment_mn; } -// On ROCM, ME and FA share the backend, and hence they share the checking -// function for fundamental limitations by the GPU kernel -// caller_is_meff is added to make the TORCH_WARN message showing the correct result -template -bool check_head_dim_size_flash(sdp_params const& params, bool debug) { #if USE_ROCM_ATTENTION - if (at::cuda::device_count() == 0) { - return false; - } - // AOTriton 0.9+ supports head_dim up to 512 - const static auto max_hdim = []() { +inline int aotriton_max_hdim() { + static const int max_hdim = []() { #if AOTRITON_VERSION_CURRENT == AOTRITON_VERSION_INT(0, 11) - // gfx11xx only support hdim <= 256 on AOTriton 0.11 + // gfx11xx only support hdim <= 256 on AOTriton 0.11/0.12 auto dprops = at::cuda::getCurrentDeviceProperties(); const c10::basic_string_view arch(dprops->gcnArchName); if (arch.starts_with("gfx11")) { @@ -165,7 +157,28 @@ bool check_head_dim_size_flash(sdp_params const& params, bool debug) { return 256; #endif }(); - const auto max_size = c10::SymInt(max_hdim); + return max_hdim; +} +#endif // USE_ROCM_ATTENTION + +// For AOTriton <= 0.11: +// On ROCM, ME and FA share the backend, and hence they share the checking +// function for fundamental limitations by the GPU kernel +// caller_is_meff is added to make the TORCH_WARN message showing the correct result +// +// FIXME: revert this reuse when removing AOTriton <= 0.11 support +// +// AOTriton 0.12 supports hdim_qk != hdim_vo, but we cannot enable this in +// check_head_dim_size_flash because it changes the backend selection logic for +// FA, which can break certain workloads that rely on the behavior of rejecting +// FA for hdim_qk != hdim_vo +template +bool check_head_dim_size_flash(sdp_params const& params, bool debug) { +#if USE_ROCM_ATTENTION + if (at::cuda::device_count() == 0) { + return false; + } + const auto max_size = c10::SymInt(aotriton_max_hdim()); #else // All head_dim sizes must be equal and less than 256 const auto max_size = c10::SymInt(256); @@ -245,9 +258,20 @@ bool check_head_dim_size_flash_nested(sdp_params const& params, bool debug) { } bool check_head_dim_size_mem_efficient(sdp_params const& params, bool debug) { +#if USE_ROCM_ATTENTION +#if AOTRITON_VERSION_CURRENT < AOTRITON_VERSION_INT(0, 12) + return check_head_dim_size_flash_nested(params, debug); +#endif +#endif const auto query_size_last = params.query.sym_size(-1); const auto value_size_last = params.value.sym_size(-1); +#ifdef USE_ROCM + bool is_half = (params.query.dtype() == at::kHalf) || + (params.query.dtype() == at::kBFloat16); + const int64_t alignment = is_half ? 8 : 4; +#else const int64_t alignment = minimum_gemm_alignment(params); +#endif if (!(query_size_last == params.key.sym_size(-1) && query_size_last % alignment == 0 && query_size_last > 0 && value_size_last % alignment == 0 && value_size_last > 0)) { @@ -266,6 +290,27 @@ bool check_head_dim_size_mem_efficient(sdp_params const& params, bool debug) { } return false; } +#if USE_ROCM_ATTENTION +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) + const auto max_size = c10::SymInt(aotriton_max_hdim()); + if (!(query_size_last <= max_size && value_size_last <= max_size)) { + if (debug) { + TORCH_WARN( + "Mem efficient attention on ROCM requires last dimension of inputs to less or equal than ", + max_size, + ". ", + "Got Query.size(-1): ", + query_size_last, + ", Key.size(-1): ", + params.key.sym_size(-1), + ", Value.size(-1): ", + params.value.sym_size(-1), + " instead. (Note this limit differs among architectures)"); + } + return false; + } +#endif // AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) +#endif // USE_ROCM_ATTENTION return true; } @@ -880,11 +925,7 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) { check_all_tensors_on_device, check_mem_efficient_hardware_support, check_tensor_shapes, -#ifdef USE_ROCM - check_head_dim_size_flash -#else check_head_dim_size_mem_efficient -#endif ); for (auto& constraint : general_constraints) { if (!constraint(params, debug)) { @@ -894,11 +935,6 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) { if (has_for_nested_inputs(params)) { constexpr auto nested_constraints = c10::array_of( -#ifndef USE_ROCM // ME and FA shares backend on ROCM and thus supports training - check_requires_grad_and_nested, -#else // Meanwhile ME on ROCM share the limits of FA about head dimensions - check_head_dim_size_flash_nested, -#endif check_batch_size_nested, check_for_seq_len_0_nested_tensor); for (auto& constraint : nested_constraints) { diff --git a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h index e40376ae0c3a7..43e9c616e35fa 100644 --- a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h +++ b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h @@ -11,6 +11,10 @@ #include #include +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) +#define AOTRITON_V2_API_FLASH_ATTN_H // Suppress the include of deprecated flash/v2.h +#endif + //////////////////////////////////////////////////////////////////////////////// // Common macros copied from cuda/mem_eff_attention/gemm_kernel_utils.h //////////////////////////////////////////////////////////////////////////////// @@ -127,8 +131,17 @@ struct LazyTensorContext { template struct LazyTensorFunctions : public LazyTensorContext { - static aotriton::TensorView acquire(void* cookie) { +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) + using HolderType = aotriton::LazyTensor; +#else + using HolderType = void; +#endif + static aotriton::TensorView acquire(HolderType* self) { +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) + auto ctx = (LazyTensorContext*)self->cookie; +#else auto ctx = (LazyTensorContext*)cookie; +#endif if (!ctx->tensor.defined()) { auto q = ctx->like_tensor; if constexpr (kRequireZeros) { @@ -141,7 +154,7 @@ struct LazyTensorFunctions : public LazyTensorContext { return mk_aotensor(ctx->tensor, ctx->tensor_name); } - static void dispose(void* cookie) { + static void dispose(HolderType* cookie) { } }; diff --git a/aten/src/ATen/native/transformers/hip/aotriton_versions.h b/aten/src/ATen/native/transformers/hip/aotriton_versions.h index 2f5d3f0e12228..b95731119e848 100644 --- a/aten/src/ATen/native/transformers/hip/aotriton_versions.h +++ b/aten/src/ATen/native/transformers/hip/aotriton_versions.h @@ -17,4 +17,10 @@ #define AOTRITON_V3_API 0 #endif +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) +#define AOTRITON_COMPACT_VARLEN_LSE 1 +#else +#define AOTRITON_COMPACT_VARLEN_LSE 1 +#endif + #endif diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip index f3816296a1d4f..c16f7d1aad233 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip @@ -257,7 +257,6 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x } hipError_t err; // TODO: Error handling - using aotriton::v2::flash::attn_fwd; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::mk_philoxtensor; @@ -270,54 +269,32 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x auto seed_output = mk_philoxtensor(use_philox_state ? seed_t.data_ptr() : nullptr); auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr() : nullptr); auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_fwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.Sm_scale = softmax_scale; - params.L = mk_aotensor<2>(M, "M"); - params.Out = mk_aotensor(output_t, "Out"); - params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = seed; - params.philox_offset1 = offset1; - params.philox_offset2 = offset2; - params.philox_seed_output = seed_output; - params.philox_offset_output = offset_output; - params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); - params.persistent_atomic_counter = persistent_counter; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - params.varlen_type = VarlenType::None; - params.window_left = window_left; - params.window_right = window_right; - err = aotriton::v3::flash::attn_fwd(params, - aotriton::v3::flash::attn_fwd_params::kVersion, - stream); -#endif - } else { - err = attn_fwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - empty_bias, - softmax_scale, - mk_aotensor<2>(M, "M"), - mk_aotensor(output_t, "Out"), - p_dropout, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); - } + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_fwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.Sm_scale = softmax_scale; + params.L = mk_aotensor<2>(M, "M"); + params.Out = mk_aotensor(output_t, "Out"); + params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = seed; + params.philox_offset1 = offset1; + params.philox_offset2 = offset2; + params.philox_seed_output = seed_output; + params.philox_offset_output = offset_output; + params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); + params.persistent_atomic_counter = persistent_counter; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + params.varlen_type = VarlenType::None; + params.window_left = window_left; + params.window_right = window_right; + err = aotriton::v3::flash::attn_fwd(params, + aotriton::v3::flash::attn_fwd_params::kVersion, + stream); // Note: These are propagated up to the return of mha_fwd(). comments // represent the assignments at that level return {out, // output @@ -350,7 +327,7 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot std::optional window_size_right, const bool return_softmax, const std::optional& gen_) { - TORCH_CHECK(!seqused_k.has_value(), "[ROCm] mha_varlen_fwd: seqused_k must be nullopt"); + bool strided_varlen = seqused_k.has_value(); const bool paged_KV = block_table_.has_value(); TORCH_CHECK(!paged_KV, "[ROCm] mha_varlen_fwd: block_table_ must be nullopt"); TORCH_CHECK(!alibi_slopes_.has_value(), "[ROCm] mha_varlen_fwd: alibi_slopes_ must be nullopt"); @@ -425,8 +402,13 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot auto opts = q.options(); +#if AOTRITON_COMPACT_VARLEN_LSE + auto softmax_lse = at::empty({num_heads, total_q}, opts.dtype(at::kFloat)); + at::Tensor M = softmax_lse; +#else auto softmax_lse = at::empty({batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat)); at::Tensor M = softmax_lse.view({batch_size * num_heads, max_seqlen_q}); +#endif at::Tensor softmax_fa_t; // Only return softmax if there's dropout to reduce compilation time if (return_softmax) { @@ -465,7 +447,6 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot if (max_seqlen_k > 0) { hipError_t err; // TODO: Error handling - using aotriton::v2::flash::attn_fwd_compact_varlen; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::mk_philoxtensor; @@ -483,60 +464,48 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot auto seed_output = use_philox_state ? mk_philoxtensor(seed_t.data_ptr()) : nullscalar; auto offset_output = use_philox_state ? mk_philoxtensor(offset_t.data_ptr()) : nullscalar; auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_fwd_params params; - params.Q = mk_aotensor(q_padded, "q"); - params.K = mk_aotensor(k_padded, "k"); - params.V = mk_aotensor(v_padded, "v"); - params.Sm_scale = softmax_scale; - params.L = mk_aotensor<2>(M, "M"); - params.Out = mk_aotensor(out_padded, "Out"); + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_fwd_params params; + params.Q = mk_aotensor(q_padded, "q"); + params.K = mk_aotensor(k_padded, "k"); + params.V = mk_aotensor(v_padded, "v"); + params.Sm_scale = softmax_scale; + params.L = mk_aotensor<2>(M, "logsumexp"); + params.Out = mk_aotensor(out_padded, "Out"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = seed; + params.philox_offset1 = offset1; + params.philox_offset2 = offset2; + params.philox_seed_output = seed_output; + params.philox_offset_output = offset_output; + params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); + params.persistent_atomic_counter = persistent_counter; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + params.varlen_type = strided_varlen ? VarlenType::StridedVarlen : VarlenType::CompactVarlen; + if (strided_varlen) { + // seqused_k holds per-batch actual kv lengths; the kernel expects cumulative + // offsets for cu_seqlens_k so it can compute seqlen_k via differencing. + // seq_strides_k carries the real memory offsets into the KV cache. + const auto& seqused = seqused_k.value(); + const int num_seqs = seqused.size(0); + at::Tensor cu_seqlens_k_from_seqused = at::zeros({num_seqs + 1}, seqused.options()); + cu_seqlens_k_from_seqused.slice(0, 1).copy_(seqused.cumsum(0)); params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = seed; - params.philox_offset1 = offset1; - params.philox_offset2 = offset2; - params.philox_seed_output = seed_output; - params.philox_offset_output = offset_output; - params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); - params.persistent_atomic_counter = persistent_counter; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - params.varlen_type = VarlenType::CompactVarlen; - params.window_left = window_left; - params.window_right = window_right; - err = aotriton::v3::flash::attn_fwd(params, - aotriton::v3::flash::attn_fwd_params::kVersion, - stream); -#endif + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k_from_seqused, "cu_seqlens_k"); + params.seq_strides_q = mk_aotensor<1>(cu_seqlens_q, "seq_strides_q"); + params.seq_strides_k = mk_aotensor<1>(cu_seqlens_k, "seq_strides_k"); } else { - err = attn_fwd_compact_varlen(mk_aotensor(q_padded, "q"), - mk_aotensor(k_padded, "k"), - mk_aotensor(v_padded, "v"), - empty_bias, - mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"), - mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - softmax_scale, - mk_aotensor<2>(M, "M"), - mk_aotensor(out_padded, "Out"), - p_dropout, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); + params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); } + params.window_left = window_left; + params.window_right = window_right; + err = aotriton::v3::flash::attn_fwd(params, + aotriton::v3::flash::attn_fwd_params::kVersion, + stream); } else { // If seqlen_k == 0, then we have an empty tensor. We need to set the output to 0. out.zero_(); @@ -678,96 +647,44 @@ mha_bwd_aot(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x hea hipError_t err; // TODO: Error handling using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - // Fused BWD does not support SWA - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_bwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.Sm_scale = softmax_scale; - params.Out = mk_aotensor(out_t, "out"); - params.DO = mk_aotensor(dout_t, "dout"); - params.DQ = mk_aotensor(dq_t, "dq"); - params.DK = mk_aotensor(dk_t, "dk"); - params.DV = mk_aotensor(dv_t, "dv"); - params.L = mk_aotensor<2>(softmax_lse_cont, "L"); - params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = mk_aoscalartensor(philox_seed); - params.philox_offset1 = mk_aoscalartensor(philox_offset); - params.philox_offset2 = 0; - // SWA in AOTriton Kernels is treated as "Generalized Causal masks" - params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; - params.window_left = window_left; - params.window_right = window_right; - params.varlen_type = VarlenType::None; -#if AOTRITON_ALWAYS_V3_API - using sdp::aotriton_adapter::mklazy_empty_like; - using sdp::aotriton_adapter::mklazy_fp32zeros; - using sdp::aotriton_adapter::LazyTensorContext; - LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; - LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; - params.D = mklazy_empty_like<2>(&lazy_delta); - params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); -#else - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - params.D = mk_aotensor<2>(delta, "delta"); -#endif - err = aotriton::v3::flash::attn_bwd(params, - aotriton::v3::flash::attn_bwd_params::kVersion, - stream); -#endif - } else if (use_fused_bwd) { - using aotriton::v2::flash::attn_bwd_fused; - using sdp::aotriton_adapter::cast_dtype; - aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype())); - err = attn_bwd_fused(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - empty_bias, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - empty_bias, // dbb - mk_aotensor<2>(softmax_lse_cont, "L"), - p_dropout, - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } else { - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - using aotriton::v2::flash::attn_bwd; - using sdp::aotriton_adapter::cast_dtype; - aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype())); - err = attn_bwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - empty_bias, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - empty_bias, // db - mk_aotensor<2>(softmax_lse_cont, "L"), - mk_aotensor<2>(delta, "delta"), - p_dropout, - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } + // Fused BWD does not support SWA + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_bwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.Sm_scale = softmax_scale; + params.Out = mk_aotensor(out_t, "out"); + params.DO = mk_aotensor(dout_t, "dout"); + params.DQ = mk_aotensor(dq_t, "dq"); + params.DK = mk_aotensor(dk_t, "dk"); + params.DV = mk_aotensor(dv_t, "dv"); + params.L = mk_aotensor<2>(softmax_lse_cont, "L"); + params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = mk_aoscalartensor(philox_seed); + params.philox_offset1 = mk_aoscalartensor(philox_offset); + params.philox_offset2 = 0; + // SWA in AOTriton Kernels is treated as "Generalized Causal masks" + params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; + params.window_left = window_left; + params.window_right = window_right; + params.varlen_type = VarlenType::None; + using sdp::aotriton_adapter::mklazy_empty_like; + using sdp::aotriton_adapter::mklazy_fp32zeros; + using sdp::aotriton_adapter::LazyTensorContext; + LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; + LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; + params.D = mklazy_empty_like<2>(&lazy_delta); + params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); + aotriton::v3::flash::attn_options options; + options.deterministic = deterministic; + err = aotriton::v3::flash::attn_bwd(params, + aotriton::v3::flash::attn_bwd_params::kVersion, + stream, + &options); return { dq, dk, dv, softmax_d }; } @@ -850,7 +767,11 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size CHECK_SHAPE(cu_seqlens_q, batch_size + 1); CHECK_SHAPE(cu_seqlens_k, batch_size + 1); +#if AOTRITON_COMPACT_VARLEN_LSE + at::Tensor softmax_lse_cont = softmax_lse.view({num_heads, total_q}).contiguous(); +#else at::Tensor softmax_lse_cont = softmax_lse.view({batch_size * num_heads, max_seqlen_q}).contiguous(); +#endif at::Tensor q_padded, k_padded, v_padded; q_padded = q.unsqueeze(0).transpose(1, 2); @@ -931,79 +852,45 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size hipError_t err; // TODO: Error handling using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_bwd_params params; - params.Q = mk_aotensor(q_padded, "q"); - params.K = mk_aotensor(k_padded, "k"); - params.V = mk_aotensor(v_padded, "v"); - params.Sm_scale = softmax_scale; - params.Out = mk_aotensor(out_t, "out"); - params.DO = mk_aotensor(dout_t, "dout"); - params.DK = mk_aotensor(dk_padded, "dk"); - params.DV = mk_aotensor(dv_padded, "dv"); - params.DQ = mk_aotensor(dq_padded, "dq"); - params.L = mk_aotensor<2>(softmax_lse_cont, "L"); - params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = mk_aoscalartensor(philox_seed); - params.philox_offset1 = mk_aoscalartensor(philox_offset); - params.philox_offset2 = 0; - // SWA in AOTriton Kernels is treated as "Generalized Causal masks" - params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; - params.varlen_type = VarlenType::CompactVarlen; - params.window_left = window_left; - params.window_right = window_right; -#if AOTRITON_ALWAYS_V3_API - using sdp::aotriton_adapter::mklazy_empty_like; - using sdp::aotriton_adapter::mklazy_fp32zeros; - using sdp::aotriton_adapter::LazyTensorContext; - LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; - LazyTensorContext lazy_dq_acc { .like_tensor = dq_padded, .tensor_name = "dq_acc" }; - params.D = mklazy_empty_like<2>(&lazy_delta); - params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); -#else - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - params.D = mk_aotensor<2>(delta, "delta"); -#endif - err = aotriton::v3::flash::attn_bwd(params, - aotriton::v3::flash::attn_bwd_params::kVersion, - stream); -#endif // AOTRITON_ALWAYS_V3_API - } else { - using aotriton::v2::flash::attn_bwd_compact_varlen; - using sdp::aotriton_adapter::cast_dtype; - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype())); - err = attn_bwd_compact_varlen(mk_aotensor(q_padded, "q"), - mk_aotensor(k_padded, "k"), - mk_aotensor(v_padded, "v"), - mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"), - mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - empty_bias, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_padded, "dq"), - mk_aotensor(dk_padded, "dk"), - mk_aotensor(dv_padded, "dv"), - empty_bias, - mk_aotensor<2>(softmax_lse_cont, "L"), - mk_aotensor<2>(delta, "delta"), - p_dropout, - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_bwd_params params; + params.Q = mk_aotensor(q_padded, "q"); + params.K = mk_aotensor(k_padded, "k"); + params.V = mk_aotensor(v_padded, "v"); + params.Sm_scale = softmax_scale; + params.Out = mk_aotensor(out_t, "out"); + params.DO = mk_aotensor(dout_t, "dout"); + params.DK = mk_aotensor(dk_padded, "dk"); + params.DV = mk_aotensor(dv_padded, "dv"); + params.DQ = mk_aotensor(dq_padded, "dq"); + params.L = mk_aotensor<2>(softmax_lse_cont, "L"); + params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = mk_aoscalartensor(philox_seed); + params.philox_offset1 = mk_aoscalartensor(philox_offset); + params.philox_offset2 = 0; + // SWA in AOTriton Kernels is treated as "Generalized Causal masks" + params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; + params.varlen_type = VarlenType::CompactVarlen; + params.window_left = window_left; + params.window_right = window_right; + using sdp::aotriton_adapter::mklazy_empty_like; + using sdp::aotriton_adapter::mklazy_fp32zeros; + using sdp::aotriton_adapter::LazyTensorContext; + LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; + LazyTensorContext lazy_dq_acc { .like_tensor = dq_padded, .tensor_name = "dq_acc" }; + params.D = mklazy_empty_like<2>(&lazy_delta); + params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); + aotriton::v3::flash::attn_options options; + options.deterministic = deterministic; + err = aotriton::v3::flash::attn_bwd(params, + aotriton::v3::flash::attn_bwd_params::kVersion, + stream, + &options); } else { // If seqlen_q == 0, then we have an empty tensor. We need to set the output to 0. dq.zero_(); diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake index d2f1e15ff1186..95b4fd94d757f 100644 --- a/cmake/External/aotriton.cmake +++ b/cmake/External/aotriton.cmake @@ -9,18 +9,14 @@ if(NOT __AOTRITON_INCLUDED) # Replaces .ci/docker/aotriton_version.txt # Note packages information may have versions skipped (due to no ABI breaks) # But they must be listed from lower version to higher version - set(__AOTRITON_VER "0.11.2b") + set(__AOTRITON_VER "0.12b") set(__AOTRITON_MANYLINUX_LIST - "manylinux_2_28" # rocm6.2 - "manylinux_2_28" # rocm6.3 "manylinux_2_28" # rocm6.4 "manylinux_2_28" # rocm7.0 "manylinux_2_28" # rocm7.1 "manylinux_2_28" # rocm7.2 ) set(__AOTRITON_ROCM_LIST - "rocm6.2" - "rocm6.3" "rocm6.4" "rocm7.0" "rocm7.1" @@ -29,31 +25,34 @@ if(NOT __AOTRITON_INCLUDED) if(DEFINED ENV{PYTORCH_AOTRITON_COMMIT}) set(__AOTRITON_CI_COMMIT "$ENV{PYTORCH_AOTRITON_COMMIT}") else() - set(__AOTRITON_CI_COMMIT "dd1b68b604b5258ee7a9f7b66ad95e7a82c18065") + set(__AOTRITON_CI_COMMIT "269036897bcee4292f4e928767df1e3dd0e3c8bd") endif() set(__AOTRITON_SHA256_LIST - "d784314849ba1911181dfc80cd845064ff6f0cdad10e2f4c53eb84a8b89245b9" # rocm6.2 - "f4b14dc111c334e967b28a1cf9ed4c63264c634dbdccbb5849aa9490022992f7" # rocm6.3 - "6b51d8479c85b902334e4f5518f404a8f5d563fd8d4732cb8b621ed4b45c2876" # rocm6.4 - "5501a0a3b300890001b6625f2a3539a7bad60f386f0a061ebe7d4ed5ca0fafb9" # rocm7.0 - "fee36beb3ea484ce18155bbafe026c577fd6705e4469e59405b260bd74b8cc10" # rocm7.1 - "cd8abf27bbb63cec45c94135e9b28745966074263a6b0555e5878ae1cb6a2349" # rocm7.2 + "e57ad080bd87fdaf7fe5bbff49ab80222be3bc8eb56a197d5781bfae8c116c33" # rocm6.4 + "f6aba1fe59312004ccd13dfda4d0a9e35457527fec18348d2e75a294a1051ef5" # rocm7.0 + "a1d731745929b61598d088eeaad31d8a82a27032aa51e3a7c831e45a99e095c2" # rocm7.1 + "5b97e8d041b160c84085961f3d3bd7b9890642b146bacb04c991aa9ad6a8dca8" # rocm7.2 ) set(__AOTRITON_IMAGE_LIST "amd-gfx90a" "amd-gfx942" "amd-gfx950" - "amd-gfx11xx" + "amd-gfx110x" + "amd-gfx115x" "amd-gfx120x" ) set(__AOTRITON_IMAGE_SHA256_LIST - "fe9f04b66bf52ac27cd025e1d89cfd04974dd3fb3ae076192f783641a4d80fdf" # amd-gfx90a - "0a7bcee19d3bb6d548732248c3234f7b92736c2ab7a7aae65294b87a7fd64c06" # amd-gfx942 - "c1ba3bfe84217fd67df3dd1f8b67c80a7f7b33d0ad4d74b41d6567036e032ace" # amd-gfx950 - "839299637fccb13fbe3e7823d57d1b2dcd0e0bed78abbcb7005ea5f4fd82b928" # amd-gfx11xx - "0a4ff324bffdac0c2fde87a8a7f70563d3c84a80ad4e8f31345f2b40a1384e95" # amd-gfx120x + "bb8bf2237b77fc503bc2967ea0d99d6ca419126c479e951ea42b712737128086" # amd-gfx90a + "f08edacf83c9ccf1c4bdcb51f1cab052d1680abea31c9e035f3f9fadb2f13ba4" # amd-gfx942 + "307a37d729cda3a2120449909e5192cd71c2badccbd37f0222786098e69c7a91" # amd-gfx950 + "c9cac7cf6f277168e1659ac2f04706f8823580b7c7e3e895f5a5503ed6bdd55f" # amd-gfx110x + "3177387a15c678b30057f4584d1fc1b8f8db56163890cb5c98f27450209f5a7b" # amd-gfx115x + "68572511ce6487a83f9014bd255bd69c8943f87d0c93bd57b2daac5fbc6c79c1" # amd-gfx120x ) - set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/") # @lint-ignore + set(__AOTRITON_BASE_URL "$ENV{PYTORCH_AOTRITON_BASE_URL}") + if(NOT __AOTRITON_BASE_URL) + set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/") # @lint-ignore + endif() set(__AOTRITON_Z "gz") # Set the default __AOTRITON_LIB path if(NOT WIN32) diff --git a/test/test_transformers.py b/test/test_transformers.py index ced9b0133e11d..ea89876790944 100644 --- a/test/test_transformers.py +++ b/test/test_transformers.py @@ -3145,7 +3145,6 @@ def test_cudnn_attention_broken_166211(self): self.assertFalse(dk.isnan().any()) self.assertFalse(dv.isnan().any()) - @skipIfRocm @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system") def test_cudnn_attention_mask_broken_177842(self): # https://github.com/pytorch/pytorch/issues/177842 @@ -3356,8 +3355,6 @@ def test_scaled_dot_product_attention_fused_kernels_packed(self, device, type: s @parametrize("type", ["nested"]) @parametrize("is_contiguous", [True, False]) def test_scaled_dot_product_attention_cudnn_nested(self, device, type: str, is_contiguous: bool): - if TEST_WITH_ROCM and type == 'nested': - self.skipTest("ROCM does not support efficient attention on nested tensors, for now") make_tensor = partial(rand_sdpa_tensor, type=type, device=device, dtype=torch.float16, packed=True) batch_size, seq_len, num_heads, head_dim = 8, 64, 16, 64 @@ -3655,7 +3652,6 @@ def compiled_func(order): reset_order = torch._C._get_sdp_priority_order() self.assertEqual(default_order, reset_order, "expected SDPA context manager to reset priority order.") - @skipIfRocm # Missing deterministic algo @unittest.skipIf(not PLATFORM_SUPPORTS_FUSED_ATTENTION, "Fused SDPA was not built for this system") @parametrize("fused_kernel", PLATFORM_SPECIFIC_SDPA) @parametrize("warn_only", [True, False]) @@ -4378,9 +4374,6 @@ def test_fused_kernels_nested_broadcasting( rand_nested_tensor = partial(rand_sdpa_tensor, type="nested", device=device, dtype=dtype) batch, num_heads, head_dim = 32, 8, 64 head_dim_v = 32 if is_efficient else head_dim - if TEST_WITH_ROCM and head_dim != head_dim_v: - self.skipTest("head_dim != head_dim_v unsupported on ROCm for now") - return seq_lens_q = (torch.randint(low=1, high=5, size=(1,)).item() if expand_q_batch else torch.randint(low=1, high=32, size=(batch,)).tolist()) @@ -4444,7 +4437,6 @@ def _broadcast(t, batch_broadcasted, num_heads_broadcasted): self.assertEqual(actual.contiguous(), math_ref.contiguous().to(dtype), atol=1.5e-3, rtol=1e-2) - @skipIfRocm(msg="Efficient Attention on ROCM does not support head_dim != head_dim_v for now.") @unittest.skipIf(not PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, "Fused SDPA was not built for this system") def test_fused_kernels_nested_broadcasting_query_dense(self, device): rand_nested_tensor = partial(rand_sdpa_tensor, type="nested", device=device, dtype=torch.float32) diff --git a/test/test_varlen_attention.py b/test/test_varlen_attention.py index dd382aea0bf05..f1283857a4d83 100644 --- a/test/test_varlen_attention.py +++ b/test/test_varlen_attention.py @@ -831,7 +831,6 @@ def test_batch_invariance( @unittest.skipIf( not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported" ) - @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support seqused_k") @decorateIf( unittest.expectedFailure, lambda params: params["backend"] != "fa2" @@ -857,6 +856,8 @@ def test_seqused_k_kv_cache( self, device, dtype, actual_kv_lens, backend, sdpa_backend=None ): if TEST_WITH_ROCM: + if sdpa_backend == "ck": + self.skipTest("CK backend does not support seqused_k") torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend) torch.manual_seed(42) @@ -961,7 +962,7 @@ def test_seqused_k_kv_cache( @unittest.skipIf( not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported" ) - @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support seqused_k") + @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support block_table") @parametrize("dtype", [torch.bfloat16, torch.float16]) @parametrize("page_size", [32, 64, 128, 256]) @parametrize("compile", [False, True]) diff --git a/torch/nn/attention/varlen.py b/torch/nn/attention/varlen.py index 96cfc332b8699..454bb82c45e2a 100644 --- a/torch/nn/attention/varlen.py +++ b/torch/nn/attention/varlen.py @@ -168,15 +168,6 @@ def _varlen_attn_fake( (num_heads, total_q), dtype=torch.float, device=query.device ) - if torch.version.hip: - preferred = torch._C._get_rocm_fa_preferred_backend() - if preferred == torch._C._ROCmFABackend.AOTriton: - # AOTriton ROCm path uses batched 3D - batch_size = cu_seq_q.size(0) - 1 - logsumexp = torch.empty( - (batch_size, num_heads, max_q), dtype=torch.float, device=query.device - ) - rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device) return output, logsumexp, rng_state @@ -416,14 +407,6 @@ def _varlen_attn_out_fake( (num_heads, total_q), dtype=torch.float, device=query.device ) - if torch.version.hip: - preferred = torch._C._get_rocm_fa_preferred_backend() - if preferred == torch._C._ROCmFABackend.AOTriton: - batch_size = cu_seq_q.size(0) - 1 - logsumexp = torch.empty( - (batch_size, num_heads, max_q), dtype=torch.float, device=query.device - ) - return logsumexp