From bcb99546b6a6b57742eb0ec1705e55981ed071be Mon Sep 17 00:00:00 2001 From: "Sv. Lockal" Date: Wed, 14 Oct 2026 15:58:21 +0800 Subject: [PATCH] Fix build with aotriton-0.14 Only the VarlenType -> VarlenBits part of the upstream commit; aotriton < 0.14 keeps the old code path. Backports: https://github.com/pytorch/pytorch/commit/9b9978943e4030e97eeee36a9968db27a3b21163 --- .../native/transformers/cuda/attention.cu | 14 +++++- .../transformers/cuda/attention_backward.cu | 14 +++++- .../transformers/hip/aotriton_adapter.h | 41 +++++++++++++++ .../transformers/hip/aotriton_versions.h | 10 ++++ .../hip/flash_attn/aot/mha_all_aot.hip | 50 +++++++++++++++++-- 5 files changed, 122 insertions(+), 7 deletions(-) diff --git a/aten/src/ATen/native/transformers/cuda/attention.cu b/aten/src/ATen/native/transformers/cuda/attention.cu index 96e3d90bd8411..537e815dbef1b 100644 --- a/aten/src/ATen/native/transformers/cuda/attention.cu +++ b/aten/src/ATen/native/transformers/cuda/attention.cu @@ -1761,8 +1761,12 @@ std::tuple _efficient_ auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); hipError_t err; // TODO: Error handling using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; using aotriton::v3::flash::WindowValue; +#if AOTRITON_VARLEN_BITS_API + using sdp::aotriton_adapter::mk_varlen_bits_packed; +#else + using aotriton::v3::flash::VarlenType; +#endif aotriton::v3::flash::attn_fwd_params params; params.Q = mk_aotensor(q_t, "q"); params.K = mk_aotensor(k_t, "k"); @@ -1791,6 +1795,13 @@ std::tuple _efficient_ if (bias.has_value()) { params.B = mk_aotensor(bias.value(), "bias"); } +#if AOTRITON_VARLEN_BITS_API + if (seqstart_q.has_value()) { + params.varlen_bits = mk_varlen_bits_packed(); + params.seqinfo_q0 = mk_aotensor<1>(seqstart_q.value(), "seqinfo_q0"); + params.seqinfo_k0 = mk_aotensor<1>(seqstart_k.value(), "seqinfo_k0"); + } +#else if (seqstart_q.has_value()) { params.varlen_type = VarlenType::CompactVarlen; params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"); @@ -1798,6 +1809,7 @@ std::tuple _efficient_ } else { params.varlen_type = VarlenType::None; } +#endif err = aotriton::v3::flash::attn_fwd(params, aotriton::v3::flash::attn_fwd_params::kVersion, stream); diff --git a/aten/src/ATen/native/transformers/cuda/attention_backward.cu b/aten/src/ATen/native/transformers/cuda/attention_backward.cu index 3b8c29c2a8052..c4756d161b434 100644 --- a/aten/src/ATen/native/transformers/cuda/attention_backward.cu +++ b/aten/src/ATen/native/transformers/cuda/attention_backward.cu @@ -651,8 +651,12 @@ _efficient_attention_backward( const auto aotriton_philox_offset = use_dropout ? philox_offset : at::zeros({}, at::dtype(at::kLong)); using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; using aotriton::v3::flash::WindowValue; +#if AOTRITON_VARLEN_BITS_API + using sdp::aotriton_adapter::mk_varlen_bits_packed; +#else + using aotriton::v3::flash::VarlenType; +#endif aotriton::v3::flash::attn_bwd_params params; params.Q = mk_aotensor(q_t, "q"); params.K = mk_aotensor(k_t, "k"); @@ -687,6 +691,13 @@ _efficient_attention_backward( 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 AOTRITON_VARLEN_BITS_API + if (cu_seqlens_q.has_value()) { + params.varlen_bits = mk_varlen_bits_packed(); + params.seqinfo_q0 = mk_aotensor<1>(cu_seqlens_q.value(), "seqinfo_q0"); + params.seqinfo_k0 = mk_aotensor<1>(cu_seqlens_k.value(), "seqinfo_k0"); + } +#else 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"); @@ -694,6 +705,7 @@ _efficient_attention_backward( } else { params.varlen_type = VarlenType::None; } +#endif aotriton::v3::flash::attn_options opts; opts.deterministic = deterministic; err = aotriton::v3::flash::attn_bwd(params, diff --git a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h index 43e9c616e35fa..78bb72aa969da 100644 --- a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h +++ b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h @@ -15,6 +15,10 @@ #define AOTRITON_V2_API_FLASH_ATTN_H // Suppress the include of deprecated flash/v2.h #endif +#if AOTRITON_VARLEN_BITS_API +#include +#endif + //////////////////////////////////////////////////////////////////////////////// // Common macros copied from cuda/mem_eff_attention/gemm_kernel_utils.h //////////////////////////////////////////////////////////////////////////////// @@ -121,6 +125,43 @@ inline aotriton::TensorView<0> mk_atomictensor(const int32_t* ptr) aotriton::DType::kInt32); } +#if AOTRITON_VARLEN_BITS_API +// AOTriton 0.14+ replaced the four-value VarlenType enum with a VarlenBits +// bit-field struct: Q and K each pick how their length is given and where +// their sequence starts, independently. Dense needs no bits (zero-init). +constexpr aotriton::v3::flash::VarlenMode mk_varlen_mode(uint32_t length, uint32_t position) +{ + aotriton::v3::flash::VarlenMode mode{}; + mode.stacked = aotriton::v3::flash::VarlenStacked::THD; + mode.length = length; + mode.position = position; + return mode; +} + +// Packed varlen: one (N+1,) cumulative array per side gives that side both its +// length (by differencing) and its start position. +constexpr aotriton::v3::flash::VarlenBits mk_varlen_bits_packed() +{ + using aotriton::v3::flash::VarlenLength; + using aotriton::v3::flash::VarlenPosition; + const auto packed = mk_varlen_mode(VarlenLength::CUMULATIVE, VarlenPosition::REUSE); + return aotriton::v3::flash::VarlenBits{.qmode = packed, .kmode = packed}; +} + +// seqused_k against a packed KV cache: K reads its length from a (N,) array of +// per-sequence counts (seqinfo_k0) and its start from a separate (N+1,) +// cumulative array (seqinfo_k1). Q stays packed. +constexpr aotriton::v3::flash::VarlenBits mk_varlen_bits_seqused_k() +{ + using aotriton::v3::flash::VarlenLength; + using aotriton::v3::flash::VarlenPosition; + return aotriton::v3::flash::VarlenBits{ + .qmode = mk_varlen_mode(VarlenLength::CUMULATIVE, VarlenPosition::REUSE), + .kmode = mk_varlen_mode(VarlenLength::INDIVIDUAL, VarlenPosition::ARRAY), + }; +} +#endif // AOTRITON_VARLEN_BITS_API + #if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 11) struct LazyTensorContext { diff --git a/aten/src/ATen/native/transformers/hip/aotriton_versions.h b/aten/src/ATen/native/transformers/hip/aotriton_versions.h index b95731119e848..10b0d4227e4da 100644 --- a/aten/src/ATen/native/transformers/hip/aotriton_versions.h +++ b/aten/src/ATen/native/transformers/hip/aotriton_versions.h @@ -23,4 +23,14 @@ #define AOTRITON_COMPACT_VARLEN_LSE 1 #endif +// AOTriton 0.14 replaced attn_{fwd,bwd}_params's VarlenType enum + cu_seqlens_q/k +// (+ seq_strides_q/k) fields with a VarlenBits bit-field struct + seqinfo_q/k0 +// (+ seqinfo_q/k1) fields. Neither spelling exists in the other version's header, +// so callers must pick one at compile time. +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 14) +#define AOTRITON_VARLEN_BITS_API 1 +#else +#define AOTRITON_VARLEN_BITS_API 0 +#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 49d7917b12852..8833d1e430389 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 @@ -270,7 +270,6 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x 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); 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"); @@ -289,7 +288,9 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x 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; +#if !AOTRITON_VARLEN_BITS_API + params.varlen_type = aotriton::v3::flash::VarlenType::None; +#endif params.window_left = window_left; params.window_right = window_right; err = aotriton::v3::flash::attn_fwd(params, @@ -327,7 +328,6 @@ 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_) { - 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"); @@ -378,6 +378,14 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot CHECK_SHAPE(v, total_k, num_heads_k, head_size_og); CHECK_SHAPE(cu_seqlens_q, batch_size + 1); CHECK_SHAPE(cu_seqlens_k, batch_size + 1); + if (seqused_k.has_value()) { + // Handed to the kernel as-is, so its layout is load-bearing. + const auto& seqused = seqused_k.value(); + TORCH_CHECK(seqused.dtype() == at::kInt, "seqused_k must have dtype int32"); + CHECK_DEVICE(seqused); + CHECK_CONTIGUOUS(seqused); + CHECK_SHAPE(seqused, batch_size); + } // AOTriton's varlen API needs input shapes be // (1, num_heads, total sequence length, head dimension) @@ -465,7 +473,12 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot 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); using aotriton::v3::flash::CausalType; +#if AOTRITON_VARLEN_BITS_API + using sdp::aotriton_adapter::mk_varlen_bits_packed; + using sdp::aotriton_adapter::mk_varlen_bits_seqused_k; +#else using aotriton::v3::flash::VarlenType; +#endif aotriton::v3::flash::attn_fwd_params params; params.Q = mk_aotensor(q_padded, "q"); params.K = mk_aotensor(k_padded, "k"); @@ -484,6 +497,18 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot 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 AOTRITON_VARLEN_BITS_API + params.seqinfo_q0 = mk_aotensor<1>(cu_seqlens_q, "seqinfo_q0"); + if (seqused_k.has_value()) { + params.varlen_bits = mk_varlen_bits_seqused_k(); + params.seqinfo_k0 = mk_aotensor<1>(seqused_k.value(), "seqinfo_k0"); + params.seqinfo_k1 = mk_aotensor<1>(cu_seqlens_k, "seqinfo_k1"); + } else { + params.varlen_bits = mk_varlen_bits_packed(); + params.seqinfo_k0 = mk_aotensor<1>(cu_seqlens_k, "seqinfo_k0"); + } +#else + const bool strided_varlen = seqused_k.has_value(); params.varlen_type = strided_varlen ? VarlenType::StridedVarlen : VarlenType::CompactVarlen; if (strided_varlen) { // seqused_k holds per-batch actual kv lengths; the kernel expects cumulative @@ -501,6 +526,7 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot 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"); } +#endif params.window_left = window_left; params.window_right = window_right; err = aotriton::v3::flash::attn_fwd(params, @@ -648,7 +674,6 @@ mha_bwd_aot(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x hea using sdp::aotriton_adapter::mk_aoscalartensor; // 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"); @@ -670,7 +695,9 @@ mha_bwd_aot(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x hea 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_VARLEN_BITS_API + params.varlen_type = aotriton::v3::flash::VarlenType::None; +#endif using sdp::aotriton_adapter::mklazy_empty_like; using sdp::aotriton_adapter::mklazy_fp32zeros; using sdp::aotriton_adapter::LazyTensorContext; @@ -851,7 +878,11 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using aotriton::v3::flash::CausalType; +#if AOTRITON_VARLEN_BITS_API + using sdp::aotriton_adapter::mk_varlen_bits_packed; +#else using aotriton::v3::flash::VarlenType; +#endif aotriton::v3::flash::attn_bwd_params params; params.Q = mk_aotensor(q_padded, "q"); params.K = mk_aotensor(k_padded, "k"); @@ -863,8 +894,13 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size params.DV = mk_aotensor(dv_padded, "dv"); params.DQ = mk_aotensor(dq_padded, "dq"); params.L = mk_aotensor<2>(softmax_lse_cont, "L"); +#if AOTRITON_VARLEN_BITS_API + params.seqinfo_q0 = mk_aotensor<1>(cu_seqlens_q, "seqinfo_q0"); + params.seqinfo_k0 = mk_aotensor<1>(cu_seqlens_k, "seqinfo_k0"); +#else 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"); +#endif 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; @@ -873,7 +909,11 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size 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; +#if AOTRITON_VARLEN_BITS_API + params.varlen_bits = mk_varlen_bits_packed(); +#else params.varlen_type = VarlenType::CompactVarlen; +#endif params.window_left = window_left; params.window_right = window_right; using sdp::aotriton_adapter::mklazy_empty_like;