Skip to content

[release/2.12] Backport Upstream PR: Bump AOTriton to 0.12b (#184288) #3290

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xinyazhang wants to merge 1 commit into
base: release/2.12
Choose a base branch
from
Open

[release/2.12] Backport Upstream PR: Bump AOTriton to 0.12b (#184288) #3290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion aten/src/ATen/Context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -505,7 +505,7 @@ at::BlasBackend Context::blasPreferredBackend() {
bool Context::ckSupported() {
#ifdef USE_ROCM
static const std::vector<std::string> supported_archs = {
"gfx90a", "gfx942", "gfx950"
"gfx942", "gfx950"
};
for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) {
if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) {
Expand Down
129 changes: 43 additions & 86 deletions aten/src/ATen/native/transformers/cuda/attention.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1563,6 +1563,10 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt> _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},
Expand Down Expand Up @@ -1591,8 +1595,6 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt> _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;
Expand All @@ -1608,92 +1610,47 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, c10::SymInt, c10::SymInt> _efficient_
auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr<int64_t>() : nullptr);
auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr<int32_t>() : 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<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
params.window_left = WindowValue::TopLeftAligned;
params.window_right = WindowValue::TopLeftAligned;
} else if (static_cast<int64_t>(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<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
params.window_left = WindowValue::TopLeftAligned;
params.window_right = WindowValue::TopLeftAligned;
} else if (static_cast<int64_t>(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
Expand Down
190 changes: 61 additions & 129 deletions aten/src/ATen/native/transformers/cuda/attention_backward.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int64_t>(sdp::CustomMaskType::NoCustomMask) == custom_mask_type) {
Expand All @@ -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<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
params.window_left = WindowValue::TopLeftAligned;
params.window_right = WindowValue::TopLeftAligned;
} else if (static_cast<int64_t>(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<int64_t>(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) {
params.window_left = WindowValue::TopLeftAligned;
params.window_right = WindowValue::TopLeftAligned;
} else if (static_cast<int64_t>(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
Expand Down
Loading