From f3dc2eada66ebb0de16a8a0fd504f6b501013f1b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 7 Jul 2026 13:30:54 +0200 Subject: [PATCH 1/3] [PyTorch/common] fused attn: opt-in real-stride plumbing to cuDNN graphs cuDNN fused attention reconstructs Q/K/V strides from the NVTE_QKV_Layout enum (generateMatrixStrides), so the enum has to encode memory geometry and Python must detect it by inspecting data_ptr, which is torch.compile-hostile. cudnn-frontend itself is stride-based, so pass the real torch strides instead. Prototype, opt-in via NVTE_FUSED_ATTN_REAL_STRIDES=1 (default off = old behavior): - new experimental C API nvte_fused_attn_set_strides(q, k, v, dO): a thread-local side channel carrying strides in cuDNN dim order [b, h, s, d]; NULL clears it. No NVTETensor/NVTEShape ABI change. - F16 arbitrary-seqlen fwd/bwd graph builders use the provided strides for the Q/K/V (and dO in bwd) set_stride instead of the enum-derived ones. Dense (non-THD, non-paged) layouts only; O and dQ/dK/dV outputs keep the enum-derived strides since TE allocates them. - FADescriptor_v1 gains a real_strides array included in operator< so cached plans built for different strides are not reused. - PyTorch extension reads q/k/v (and dO) .strides() from the incoming at::Tensors, permutes them from bshd/sbhd/bhsd torch order to [b, h, s, d], sets the side channel around the nvte_fused_attn_fwd/bwd calls and clears it afterwards. Verified on sm89 (bf16, b=2 s=128 h=8 d=64, no_mask, deterministic bwd): packed bs3hd/sbh3d views passed with the *separate* layout enum (bshd_bshd_bshd / sbhd_sbhd_sbhd) are bit-exact vs contiguous baselines in forward and backward with the flag on, and wrong with the flag off (enum strides mismatch memory), demonstrating the enum no longer needs to encode memory geometry. test_dot_product_attention passes with the flag off and on. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../common/fused_attn/fused_attn.cpp | 37 +++++++ .../fused_attn_f16_arbitrary_seqlen.cu | 60 +++++++++++- transformer_engine/common/fused_attn/utils.h | 23 ++++- .../include/transformer_engine/fused_attn.h | 22 +++++ .../pytorch/csrc/extensions/attention.cpp | 96 +++++++++++++++++++ 5 files changed, 232 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..e375a3c9a4 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" @@ -94,8 +96,43 @@ std::string to_string(NVTE_QKV_Format format) { } } +namespace fused_attn { + +RealStrideOverride &GetRealStrideOverride() { + static thread_local RealStrideOverride instance; + return instance; +} + +} // namespace fused_attn + } // namespace transformer_engine +// Prototype: set/clear the thread-local real-stride override for Q/K/V (+dO) +void nvte_fused_attn_set_strides(const int64_t *q_strides, const int64_t *k_strides, + const int64_t *v_strides, const int64_t *do_strides) { + NVTE_API_CALL(nvte_fused_attn_set_strides); + using namespace transformer_engine::fused_attn; + auto &override = GetRealStrideOverride(); + if (q_strides != nullptr && k_strides != nullptr && v_strides != nullptr) { + std::copy(q_strides, q_strides + 4, override.q.begin()); + std::copy(k_strides, k_strides + 4, override.k.begin()); + std::copy(v_strides, v_strides + 4, override.v.begin()); + override.has_qkv = true; + } else { + override.has_qkv = false; + override.q.fill(0); + override.k.fill(0); + override.v.fill(0); + } + if (do_strides != nullptr) { + std::copy(do_strides, do_strides + 4, override.dO.begin()); + override.has_do = true; + } else { + override.has_do = false; + override.dO.fill(0); + } +} + // map NVTE_QKV_Layout to NVTE_QKV_Layout_Group NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { switch (qkv_layout) { 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..eb386ca75d 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 @@ -94,6 +94,11 @@ void fused_attn_arbitrary_seqlen_fwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + // Prototype: real-stride override for dense (non-THD, non-paged) Q/K/V + const auto &stride_override = GetRealStrideOverride(); + const bool use_real_strides = + stride_override.has_qkv && !is_paged_kv && !is_ragged_q && !is_ragged_kv; + // 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) { @@ -153,6 +158,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cudnn_frontend::DataType_t::NOT_SET, return_max_logit, }; + if (use_real_strides) { + for (int i = 0; i < 4; ++i) { + descriptor.real_strides[i] = stride_override.q[i]; + descriptor.real_strides[4 + i] = stride_override.k[i]; + descriptor.real_strides[8 + i] = stride_override.v[i]; + } + } namespace fe = cudnn_frontend; using graph_and_tensors = @@ -219,6 +231,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); } + if (use_real_strides) { + // use the real tensor strides instead of the enum-derived ones + q_stride.assign(stride_override.q.begin(), stride_override.q.end()); + k_stride.assign(stride_override.k.begin(), stride_override.k.end()); + v_stride.assign(stride_override.v.begin(), stride_override.v.end()); + } Q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Q") @@ -598,6 +616,13 @@ void fused_attn_arbitrary_seqlen_bwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } + // Prototype: real-stride override for dense (non-THD, non-paged) Q/K/V inputs and dO. + // Outputs dQ/dK/dV keep the enum-derived strides (TE allocates them). + const auto &stride_override = GetRealStrideOverride(); + const bool is_dense = !is_paged_kv && !is_ragged_q && !is_ragged_kv; + const bool use_real_strides = stride_override.has_qkv && is_dense; + const bool use_real_do_strides = stride_override.has_do && is_dense; + // 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) { @@ -656,6 +681,18 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cudnn_frontend::DataType_t::NOT_SET, false, }; + if (use_real_strides) { + for (int i = 0; i < 4; ++i) { + descriptor.real_strides[i] = stride_override.q[i]; + descriptor.real_strides[4 + i] = stride_override.k[i]; + descriptor.real_strides[8 + i] = stride_override.v[i]; + } + } + if (use_real_do_strides) { + for (int i = 0; i < 4; ++i) { + descriptor.real_strides[12 + i] = stride_override.dO[i]; + } + } namespace fe = cudnn_frontend; using graph_and_tensors = @@ -722,6 +759,21 @@ void fused_attn_arbitrary_seqlen_bwd_impl( generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_O_Matrix); + // dQ/dK/dV outputs and O keep the enum-derived strides; the real-stride override + // only applies to the Q/K/V and dO inputs. + std::vector dq_stride(q_stride); + std::vector dk_stride(k_stride); + std::vector dv_stride(v_stride); + std::vector do_stride(o_stride); + if (use_real_strides) { + q_stride.assign(stride_override.q.begin(), stride_override.q.end()); + k_stride.assign(stride_override.k.begin(), stride_override.k.end()); + v_stride.assign(stride_override.v.begin(), stride_override.v.end()); + } + if (use_real_do_strides) { + do_stride.assign(stride_override.dO.begin(), stride_override.dO.end()); + } + q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Q") .set_dim({b, h, s_q, d_qk}) @@ -741,7 +793,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( dO = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("dO") .set_dim({b, h, s_q, d_v}) - .set_stride(o_stride)); + .set_stride(do_stride)); if (is_ragged_q) { offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("offset_q") @@ -893,9 +945,9 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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); + dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(dq_stride); + dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(dk_stride); + dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(dv_stride); if (is_ragged_q) { dQ->set_ragged_offset(offset_q); } diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 41656062a4..d9770735e9 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -312,6 +313,10 @@ struct FADescriptor_v1 { cudnn_frontend::DataType_t do_tensor_type; cudnn_frontend::DataType_t dqkv_tensor_type; bool return_max_logit; + // Real (torch) strides for Q, K, V, dO in cuDNN dim order [b, h, s, d], concatenated + // ([0:4]=Q, [4:8]=K, [8:12]=V, [12:16]=dO). All zeros when the enum-derived strides + // are used. Part of the cache key so plans built for different strides are not reused. + std::array real_strides = {}; 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, @@ -320,7 +325,7 @@ struct FADescriptor_v1 { 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) < + dqkv_tensor_type, return_max_logit, real_strides) < 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, @@ -329,10 +334,24 @@ struct FADescriptor_v1 { 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); + rhs.dqkv_tensor_type, rhs.return_max_logit, rhs.real_strides); } }; +// Prototype: thread-local override of Q/K/V (+dO) strides for the F16 arbitrary-seqlen +// backend, set via nvte_fused_attn_set_strides. Strides are in cuDNN dim order +// [b, h, s, d], in units of elements. +struct RealStrideOverride { + bool has_qkv = false; + bool has_do = false; + std::array q{}; + std::array k{}; + std::array v{}; + std::array dO{}; +}; + +RealStrideOverride &GetRealStrideOverride(); + __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 41e4b136bd..1aaa5ee579 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -388,6 +388,28 @@ void nvte_populate_rng_state_async(NVTETensor rng_state_dst, const NVTETensor se size_t q_max_seqlen, size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend, cudaStream_t stream); +/*! \brief Provide real tensor strides for Q/K/V (and optionally dO) to subsequent + * nvte_fused_attn_fwd / nvte_fused_attn_bwd calls on the calling thread. + * + * \warning This API is **experimental** and subject to change. + * + * Prototype side channel: when set, the F16 arbitrary-seqlen backend uses these strides + * for the cuDNN graph tensors Q/K/V (and dO in backward) instead of reconstructing + * strides from the NVTE_QKV_Layout enum. Only applies to dense (non-THD, non-paged) + * layouts; ignored otherwise. The override stays active for the calling thread until + * cleared by passing NULL pointers. + * + * \param[in] q_strides Q strides, 4 elements in cuDNN dim order [b, h, s, d], + * in units of elements. NULL (together with k/v) clears the + * Q/K/V override. + * \param[in] k_strides K strides, same convention as q_strides. + * \param[in] v_strides V strides, same convention as q_strides. + * \param[in] do_strides dO strides (backward only), same convention. NULL clears + * the dO override. + */ +void nvte_fused_attn_set_strides(const int64_t *q_strides, const int64_t *k_strides, + const int64_t *v_strides, const int64_t *do_strides); + /*! \brief Get KV format for a given QKV layout. * * \warning This API is **experimental** and subject to change. diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index eb8813d4a0..4763b58553 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -4,6 +4,8 @@ * See LICENSE for license information. ************************************************************************/ +#include + #include "../extensions.h" #include "common.h" #include "pybind.h" @@ -12,6 +14,54 @@ namespace { constexpr int block_size = 512; +// Prototype (opt-in via NVTE_FUSED_ATTN_REAL_STRIDES=1): pass the real torch strides of +// Q/K/V (and dO in backward) down to the cuDNN fused-attention graph instead of relying +// on strides reconstructed from the NVTE_QKV_Layout enum. +bool real_strides_enabled() { + const char *env = std::getenv("NVTE_FUSED_ATTN_REAL_STRIDES"); + return env != nullptr && env[0] == '1'; +} + +// Cast a py::handle to at::Tensor if it wraps a plain torch tensor. +bool get_plain_torch_tensor(pybind11::handle handle, at::Tensor &out) { + try { + out = handle.cast(); + return true; + } catch (const pybind11::cast_error &) { + return false; + } +} + +// Map a 4-D torch tensor's strides to cuDNN dim order [b, h, s, d] given its format. +// Returns false for unsupported formats (e.g. THD) or non-4D tensors. +bool to_bhsd_strides(const at::Tensor &t, NVTE_QKV_Format format, int64_t *out) { + if (t.dim() != 4) { + return false; + } + switch (format) { + case NVTE_QKV_Format::NVTE_BSHD: // torch [b, s, h, d] + out[0] = t.stride(0); + out[1] = t.stride(2); + out[2] = t.stride(1); + out[3] = t.stride(3); + return true; + case NVTE_QKV_Format::NVTE_SBHD: // torch [s, b, h, d] + out[0] = t.stride(1); + out[1] = t.stride(2); + out[2] = t.stride(0); + out[3] = t.stride(3); + return true; + case NVTE_QKV_Format::NVTE_BHSD: // torch [b, h, s, d] + out[0] = t.stride(0); + out[1] = t.stride(1); + out[2] = t.stride(2); + out[3] = t.stride(3); + return true; + default: + return false; + } +} + // fast zero-fills of tensors void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &start_index) { std::vector shape = transformer_engine::pytorch::convertShape(self.shape()); @@ -243,6 +293,23 @@ std::vector fused_attn_fwd( // create workspace TensorWrapper workspace; + // Prototype: pass real torch strides of Q/K/V through to the cuDNN graph + bool real_strides_set = false; + if (real_strides_enabled()) { + at::Tensor q_torch, k_torch, v_torch; + if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && + get_plain_torch_tensor(V, v_torch)) { + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + int64_t q_strides[4], k_strides[4], v_strides[4]; + if (to_bhsd_strides(q_torch, q_format, q_strides) && + to_bhsd_strides(k_torch, kv_format, k_strides) && + to_bhsd_strides(v_torch, kv_format, v_strides)) { + nvte_fused_attn_set_strides(q_strides, k_strides, v_strides, nullptr); + real_strides_set = true; + } + } + } + // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd( @@ -312,6 +379,11 @@ std::vector fused_attn_fwd( window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); }); + // clear the real-stride override + if (real_strides_set) { + nvte_fused_attn_set_strides(nullptr, nullptr, nullptr, nullptr); + } + // destroy tensor wrappers, but not allocated memory nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); @@ -571,6 +643,25 @@ std::vector fused_attn_bwd( // create workspace TensorWrapper workspace; + // Prototype: pass real torch strides of Q/K/V (and dO) through to the cuDNN graph + bool real_strides_set = false; + if (real_strides_enabled()) { + at::Tensor q_torch, k_torch, v_torch, do_torch; + if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && + get_plain_torch_tensor(V, v_torch)) { + int64_t q_strides[4], k_strides[4], v_strides[4], do_strides[4]; + if (to_bhsd_strides(q_torch, q_format, q_strides) && + to_bhsd_strides(k_torch, kv_format, k_strides) && + to_bhsd_strides(v_torch, kv_format, v_strides)) { + const bool have_do_strides = get_plain_torch_tensor(dO, do_torch) && + to_bhsd_strides(do_torch, do_format, do_strides); + nvte_fused_attn_set_strides(q_strides, k_strides, v_strides, + have_do_strides ? do_strides : nullptr); + real_strides_set = true; + } + } + } + // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd( @@ -602,6 +693,11 @@ std::vector fused_attn_bwd( at::cuda::getCurrentCUDAStream()); }); + // clear the real-stride override + if (real_strides_set) { + nvte_fused_attn_set_strides(nullptr, nullptr, nullptr, nullptr); + } + // destroy tensor wrappers nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); From 4eaa4f6b717723a47d6635f3f6d374f1440c4b48 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 7 Jul 2026 15:15:48 +0200 Subject: [PATCH 2/3] [PyTorch/common] fused attn: replace TLS stride side channel with explicit v2 API Convert the real-stride prototype plumbing to a production shape: instead of the thread-local nvte_fused_attn_set_strides side channel, thread the strides through an explicit, additive C API. - new C struct NVTEQKVStrides {q, k, v, d_o}: each member points to 4 int64 element-strides in cuDNN dim order [b, h, s, d], or is NULL to derive the strides from the NVTE_QKV_Layout enum as before. - new exported functions nvte_fused_attn_fwd_v2 / nvte_fused_attn_bwd_v2: identical to the v1 entry points plus an NVTEQKVStrides parameter. The old nvte_fused_attn_fwd/bwd become thin wrappers calling _v2 with all-NULL strides, so existing callers (JAX, external) are untouched and behavior is unchanged. - strides-first core: the parameter is threaded explicitly through fused_attn_arbitrary_seqlen_fwd/bwd into the F16 arbitrary-seqlen fwd/bwd_impl graph builders (no globals). Provided strides are used for the Q/K/V (and dO in bwd) set_stride on dense (non-THD, non-paged) layouts; otherwise the enum-derived strides are used exactly as before. Outputs (O, dQ/dK/dV) keep enum-derived strides since TE allocates them. FP8 and max-512 backends ignore the parameter (documented in the header). - FADescriptor_v1 keeps the real_strides cache-key extension, now populated from the threaded parameter; the RealStrideOverride TLS machinery and nvte_fused_attn_set_strides are removed entirely. Since the strides are an explicit argument, both phases of the two-phase (workspace-size + execute) call sequence trivially receive the same values. - PyTorch extension builds NVTEQKVStrides from the incoming at::Tensor strides (bshd/sbhd/bhsd -> [b, h, s, d] permutation as before) and calls _v2 directly; still gated by NVTE_FUSED_ATTN_REAL_STRIDES=1 (default off). - add tests/pytorch/attention/test_fused_attn_real_strides.py covering flag-on == flag-off for contiguous inputs, bit-exact packed bs3hd/sbh3d strided views declared with separate layout enums, and the unchanged packed-enum path. Verified on sm89: full prototype matrix bit-exact (packed views declared separate == contiguous baselines, fwd+bwd; negative control still wrong with the flag off), test_dot_product_attention 48 passed/48 skipped with the flag off and on, and old C symbols still exported. Co-Authored-By: Claude Fable 5 Signed-off-by: Pawel Gadzinski --- .../attention/test_fused_attn_real_strides.py | 159 +++++++++++++++++ .../common/fused_attn/fused_attn.cpp | 164 ++++++++++-------- .../fused_attn_f16_arbitrary_seqlen.cu | 105 +++++------ .../fused_attn_f16_arbitrary_seqlen.h | 16 +- transformer_engine/common/fused_attn/utils.h | 14 -- .../include/transformer_engine/fused_attn.h | 93 ++++++++-- .../pytorch/csrc/extensions/attention.cpp | 60 +++---- 7 files changed, 411 insertions(+), 200 deletions(-) create mode 100644 tests/pytorch/attention/test_fused_attn_real_strides.py diff --git a/tests/pytorch/attention/test_fused_attn_real_strides.py b/tests/pytorch/attention/test_fused_attn_real_strides.py new file mode 100644 index 0000000000..f6589532b4 --- /dev/null +++ b/tests/pytorch/attention/test_fused_attn_real_strides.py @@ -0,0 +1,159 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for NVTE_FUSED_ATTN_REAL_STRIDES (nvte_fused_attn_fwd/bwd_v2). + +With the flag on, the PyTorch extension passes the real torch strides of +Q/K/V (and dO) to the cuDNN fused-attention graph instead of strides +reconstructed from the NVTE_QKV_Layout enum. Strided views into a packed +QKV buffer then compute correctly even when declared with the plain +*separate* layout enum -- the enum no longer needs to encode memory +geometry. +""" + +import os + +import pytest +import torch + +import transformer_engine.pytorch # noqa: F401 (loads libtransformer_engine.so) +import transformer_engine_torch as tex +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + fused_attn_bwd, + fused_attn_fwd, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") + +_BACKEND = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen +_B, _S, _H, _D = 2, 128, 8, 64 +_DTYPE = torch.bfloat16 + + +@pytest.fixture +def real_strides_flag(monkeypatch): + def set_flag(on: bool): + if on: + monkeypatch.setenv("NVTE_FUSED_ATTN_REAL_STRIDES", "1") + else: + monkeypatch.delenv("NVTE_FUSED_ATTN_REAL_STRIDES", raising=False) + + yield set_flag + monkeypatch.delenv("NVTE_FUSED_ATTN_REAL_STRIDES", raising=False) + + +def _cu_seqlens(): + return torch.arange(0, (_B + 1) * _S, _S, dtype=torch.int32, device="cuda") + + +def _run(q, k, v, d_o, qkv_layout, fmt): + """One fwd+bwd through the F16 arbitrary-seqlen backend; returns (out, dq, dk, dv).""" + cu = _cu_seqlens() + out, aux = fused_attn_fwd( + True, _S, _S, cu, cu, q, k, v, _DTYPE, _BACKEND, + dropout=0.0, qkv_layout=qkv_layout, o_format=fmt, + attn_bias_type="no_bias", attn_mask_type="no_mask", + ) + dq, dk, dv, _, _ = fused_attn_bwd( + _S, _S, cu, cu, q, k, v, out, d_o, _DTYPE, aux, _BACKEND, + dropout=0.0, qkv_layout=qkv_layout, o_format=fmt, do_format=fmt, + dqkv_layout=qkv_layout, attn_bias_type="no_bias", attn_mask_type="no_mask", + deterministic=True, + ) + return out, dq, dk, dv + + +def _assert_bit_exact(result, reference): + for name, x, y in zip(("out", "dq", "dk", "dv"), result, reference): + assert torch.equal(x.contiguous(), y.contiguous()), f"{name} differs" + + +def _backend_supported(): + try: + q = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + _run(q, q.clone(), q.clone(), q.clone(), "bshd_bshd_bshd", "bshd") + return True + except Exception: + return False + + +requires_backend = pytest.mark.skipif( + not (torch.cuda.is_available() and _backend_supported()), + reason="F16_arbitrary_seqlen fused attention backend is not supported on this device", +) + + +@requires_backend +@pytest.mark.parametrize("fmt", ["bshd", "sbhd"]) +def test_contiguous_separate_matches_enum_path(real_strides_flag, fmt): + """Contiguous separate q/k/v: real strides == enum strides, so flag on == flag off.""" + torch.manual_seed(0) + shape = (_B, _S, _H, _D) if fmt == "bshd" else (_S, _B, _H, _D) + q = torch.randn(*shape, dtype=_DTYPE, device="cuda") + k, v, d_o = q.clone(), q.clone(), torch.randn(*shape, dtype=_DTYPE, device="cuda") + layout = f"{fmt}_{fmt}_{fmt}" + + real_strides_flag(False) + reference = _run(q, k, v, d_o, layout, fmt) + real_strides_flag(True) + result = _run(q, k, v, d_o, layout, fmt) + _assert_bit_exact(result, reference) + + +@requires_backend +def test_packed_bs3hd_views_declared_separate(real_strides_flag): + """Headline: strided views into a packed [b,s,3,h,d] buffer, declared with the + plain separate layout enum, are bit-exact vs the contiguous baseline when real + strides are passed.""" + torch.manual_seed(0) + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + d_o = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + + real_strides_flag(False) + reference = _run( + qkv[:, :, 0].contiguous(), qkv[:, :, 1].contiguous(), qkv[:, :, 2].contiguous(), + d_o, "bshd_bshd_bshd", "bshd", + ) + + real_strides_flag(True) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + assert not q.is_contiguous() + result = _run(q, k, v, d_o, "bshd_bshd_bshd", "bshd") + _assert_bit_exact(result, reference) + + +@requires_backend +def test_packed_sbh3d_views_declared_separate(real_strides_flag): + """Same as above for the sbh3d interleave (packing at dim -2, sbhd format).""" + torch.manual_seed(0) + qkv = torch.randn(_S, _B, _H, 3, _D, dtype=_DTYPE, device="cuda") + d_o = torch.randn(_S, _B, _H, _D, dtype=_DTYPE, device="cuda") + + real_strides_flag(False) + reference = _run( + qkv[:, :, :, 0].contiguous(), qkv[:, :, :, 1].contiguous(), + qkv[:, :, :, 2].contiguous(), d_o, "sbhd_sbhd_sbhd", "sbhd", + ) + + real_strides_flag(True) + q, k, v = qkv[:, :, :, 0], qkv[:, :, :, 1], qkv[:, :, :, 2] + assert not q.is_contiguous() + result = _run(q, k, v, d_o, "sbhd_sbhd_sbhd", "sbhd") + _assert_bit_exact(result, reference) + + +@requires_backend +def test_packed_views_with_packed_enum_unchanged(real_strides_flag): + """Regression: the historical path (packed views with the matching packed enum) + is unchanged with the flag on.""" + torch.manual_seed(0) + qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") + d_o = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + real_strides_flag(False) + reference = _run(q, k, v, d_o, "bs3hd", "bshd") + real_strides_flag(True) + result = _run(q, k, v, d_o, "bs3hd", "bshd") + _assert_bit_exact(result, reference) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index e375a3c9a4..b1ad326d62 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -6,8 +6,6 @@ #include "transformer_engine/fused_attn.h" -#include - #include "../common.h" #include "../cudnn_utils.h" #include "../util/cuda_runtime.h" @@ -96,43 +94,8 @@ std::string to_string(NVTE_QKV_Format format) { } } -namespace fused_attn { - -RealStrideOverride &GetRealStrideOverride() { - static thread_local RealStrideOverride instance; - return instance; -} - -} // namespace fused_attn - } // namespace transformer_engine -// Prototype: set/clear the thread-local real-stride override for Q/K/V (+dO) -void nvte_fused_attn_set_strides(const int64_t *q_strides, const int64_t *k_strides, - const int64_t *v_strides, const int64_t *do_strides) { - NVTE_API_CALL(nvte_fused_attn_set_strides); - using namespace transformer_engine::fused_attn; - auto &override = GetRealStrideOverride(); - if (q_strides != nullptr && k_strides != nullptr && v_strides != nullptr) { - std::copy(q_strides, q_strides + 4, override.q.begin()); - std::copy(k_strides, k_strides + 4, override.k.begin()); - std::copy(v_strides, v_strides + 4, override.v.begin()); - override.has_qkv = true; - } else { - override.has_qkv = false; - override.q.fill(0); - override.k.fill(0); - override.v.fill(0); - } - if (do_strides != nullptr) { - std::copy(do_strides, do_strides + 4, override.dO.begin()); - override.has_do = true; - } else { - override.has_do = false; - override.dO.fill(0); - } -} - // map NVTE_QKV_Layout to NVTE_QKV_Layout_Group NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { switch (qkv_layout) { @@ -566,22 +529,20 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return backend; } -// 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); +// NVTE fused attention FWD with separate Q, K and V and explicit strides +void nvte_fused_attn_fwd_v2( + 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, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream) { + 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); @@ -659,10 +620,11 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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); + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, qkv_strides, + 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, @@ -674,23 +636,47 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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); + nvte_fused_attn_fwd_v2(Q, K, V, Bias, SoftmaxOffset, S, O, Aux_CTX_Tensors, cu_seqlens_q, + cu_seqlens_kv, cu_seqlens_q_padded, cu_seqlens_kv_padded, page_table_k, + page_table_v, rng_state, max_seqlen_q, max_seqlen_kv, is_training, + return_max_logit, cuda_graph, 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, + NVTEQKVStrides{nullptr, nullptr, nullptr, nullptr}, workspace, stream); +} + +// NVTE fused attention BWD with separate Q, K and V and explicit strides +void nvte_fused_attn_bwd_v2( + 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, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream) { + 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); @@ -749,11 +735,11 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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); + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_strides, + 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++]); @@ -779,6 +765,32 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso } } +// 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); + nvte_fused_attn_bwd_v2(Q, K, V, O, dO, S, dP, Aux_CTX_Tensors, dQ, dK, dV, dBias, dSoftmaxOffset, + cu_seqlens_q, cu_seqlens_kv, cu_seqlens_q_padded, cu_seqlens_kv_padded, + max_seqlen_q, max_seqlen_kv, 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, cuda_graph, + NVTEQKVStrides{nullptr, nullptr, nullptr, nullptr}, workspace, stream); +} + 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 eb386ca75d..9d542bc022 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 @@ -55,12 +55,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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, - 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) { + int64_t window_size_right, bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, + 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) { using namespace transformer_engine; bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -94,10 +95,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } - // Prototype: real-stride override for dense (non-THD, non-paged) Q/K/V - const auto &stride_override = GetRealStrideOverride(); + // Real strides provided by the caller replace the enum-derived Q/K/V strides for + // dense (non-THD, non-paged) layouts. + const bool has_real_qkv_strides = + qkv_strides.q != nullptr && qkv_strides.k != nullptr && qkv_strides.v != nullptr; const bool use_real_strides = - stride_override.has_qkv && !is_paged_kv && !is_ragged_q && !is_ragged_kv; + has_real_qkv_strides && !is_paged_kv && !is_ragged_q && !is_ragged_kv; // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; @@ -160,9 +163,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( }; if (use_real_strides) { for (int i = 0; i < 4; ++i) { - descriptor.real_strides[i] = stride_override.q[i]; - descriptor.real_strides[4 + i] = stride_override.k[i]; - descriptor.real_strides[8 + i] = stride_override.v[i]; + descriptor.real_strides[i] = qkv_strides.q[i]; + descriptor.real_strides[4 + i] = qkv_strides.k[i]; + descriptor.real_strides[8 + i] = qkv_strides.v[i]; } } @@ -233,9 +236,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } if (use_real_strides) { // use the real tensor strides instead of the enum-derived ones - q_stride.assign(stride_override.q.begin(), stride_override.q.end()); - k_stride.assign(stride_override.k.begin(), stride_override.k.end()); - v_stride.assign(stride_override.v.begin(), stride_override.v.end()); + q_stride.assign(qkv_strides.q, qkv_strides.q + 4); + k_stride.assign(qkv_strides.k, qkv_strides.k + 4); + v_stride.assign(qkv_strides.v, qkv_strides.v + 4); } Q = mha_graph->tensor(fe::graph::Tensor_attributes() @@ -576,10 +579,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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, + bool bottom_right_diagonal, bool deterministic, NVTEQKVStrides qkv_strides, 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) { @@ -616,12 +619,13 @@ void fused_attn_arbitrary_seqlen_bwd_impl( NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } - // Prototype: real-stride override for dense (non-THD, non-paged) Q/K/V inputs and dO. - // Outputs dQ/dK/dV keep the enum-derived strides (TE allocates them). - const auto &stride_override = GetRealStrideOverride(); + // Real strides provided by the caller replace the enum-derived strides of the Q/K/V and + // dO inputs for dense (non-THD, non-paged) layouts. Outputs dQ/dK/dV keep the + // enum-derived strides (TE allocates them). const bool is_dense = !is_paged_kv && !is_ragged_q && !is_ragged_kv; - const bool use_real_strides = stride_override.has_qkv && is_dense; - const bool use_real_do_strides = stride_override.has_do && is_dense; + const bool use_real_strides = + qkv_strides.q != nullptr && qkv_strides.k != nullptr && qkv_strides.v != nullptr && is_dense; + const bool use_real_do_strides = qkv_strides.d_o != nullptr && is_dense; // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; @@ -683,14 +687,14 @@ void fused_attn_arbitrary_seqlen_bwd_impl( }; if (use_real_strides) { for (int i = 0; i < 4; ++i) { - descriptor.real_strides[i] = stride_override.q[i]; - descriptor.real_strides[4 + i] = stride_override.k[i]; - descriptor.real_strides[8 + i] = stride_override.v[i]; + descriptor.real_strides[i] = qkv_strides.q[i]; + descriptor.real_strides[4 + i] = qkv_strides.k[i]; + descriptor.real_strides[8 + i] = qkv_strides.v[i]; } } if (use_real_do_strides) { for (int i = 0; i < 4; ++i) { - descriptor.real_strides[12 + i] = stride_override.dO[i]; + descriptor.real_strides[12 + i] = qkv_strides.d_o[i]; } } @@ -766,12 +770,12 @@ void fused_attn_arbitrary_seqlen_bwd_impl( std::vector dv_stride(v_stride); std::vector do_stride(o_stride); if (use_real_strides) { - q_stride.assign(stride_override.q.begin(), stride_override.q.end()); - k_stride.assign(stride_override.k.begin(), stride_override.k.end()); - v_stride.assign(stride_override.v.begin(), stride_override.v.end()); + q_stride.assign(qkv_strides.q, qkv_strides.q + 4); + k_stride.assign(qkv_strides.k, qkv_strides.k + 4); + v_stride.assign(qkv_strides.v, qkv_strides.v + 4); } if (use_real_do_strides) { - do_stride.assign(stride_override.dO.begin(), stride_override.dO.end()); + do_stride.assign(qkv_strides.d_o, qkv_strides.d_o + 4); } q = mha_graph->tensor(fe::graph::Tensor_attributes() @@ -1131,12 +1135,12 @@ void fused_attn_arbitrary_seqlen_fwd( 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) { + bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, 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 auto QKV_type = input_Q->data.dtype; @@ -1266,11 +1270,11 @@ void fused_attn_arbitrary_seqlen_fwd( 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), - workspace->data.dptr, &workspace_size, stream, handle); + mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + qkv_strides, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, + devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, + devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1294,8 +1298,8 @@ void fused_attn_arbitrary_seqlen_bwd( 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, + bool deterministic, NVTEQKVStrides qkv_strides, 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, @@ -1365,11 +1369,12 @@ void fused_attn_arbitrary_seqlen_bwd( 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); + window_size_left, window_size_right, bottom_right_diagonal, deterministic, qkv_strides, + 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); if (workspace_size > 0) { if (workspace->data.dptr == 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 8f79b5bb4a..2a9414d5c2 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 @@ -25,12 +25,12 @@ void fused_attn_arbitrary_seqlen_fwd( 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); + bool bottom_right_diagonal, NVTEQKVStrides qkv_strides, 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, @@ -39,8 +39,8 @@ void fused_attn_arbitrary_seqlen_bwd( 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, + bool deterministic, NVTEQKVStrides qkv_strides, 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, diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index d9770735e9..3719616f41 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -338,20 +338,6 @@ struct FADescriptor_v1 { } }; -// Prototype: thread-local override of Q/K/V (+dO) strides for the F16 arbitrary-seqlen -// backend, set via nvte_fused_attn_set_strides. Strides are in cuDNN dim order -// [b, h, s, d], in units of elements. -struct RealStrideOverride { - bool has_qkv = false; - bool has_do = false; - std::array q{}; - std::array k{}; - std::array v{}; - std::array dO{}; -}; - -RealStrideOverride &GetRealStrideOverride(); - __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 1aaa5ee579..6f3acb4984 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -388,27 +388,84 @@ void nvte_populate_rng_state_async(NVTETensor rng_state_dst, const NVTETensor se size_t q_max_seqlen, size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend, cudaStream_t stream); -/*! \brief Provide real tensor strides for Q/K/V (and optionally dO) to subsequent - * nvte_fused_attn_fwd / nvte_fused_attn_bwd calls on the calling thread. +/*! \struct NVTEQKVStrides + * \brief Real memory strides of the Q, K, V (and dO in backward) tensors. * * \warning This API is **experimental** and subject to change. * - * Prototype side channel: when set, the F16 arbitrary-seqlen backend uses these strides - * for the cuDNN graph tensors Q/K/V (and dO in backward) instead of reconstructing - * strides from the NVTE_QKV_Layout enum. Only applies to dense (non-THD, non-paged) - * layouts; ignored otherwise. The override stays active for the calling thread until - * cleared by passing NULL pointers. - * - * \param[in] q_strides Q strides, 4 elements in cuDNN dim order [b, h, s, d], - * in units of elements. NULL (together with k/v) clears the - * Q/K/V override. - * \param[in] k_strides K strides, same convention as q_strides. - * \param[in] v_strides V strides, same convention as q_strides. - * \param[in] do_strides dO strides (backward only), same convention. NULL clears - * the dO override. - */ -void nvte_fused_attn_set_strides(const int64_t *q_strides, const int64_t *k_strides, - const int64_t *v_strides, const int64_t *do_strides); + * Each member points to 4 int64 strides in cuDNN dimension order [b, h, s, d], in units + * of elements, or is NULL, in which case the strides of the corresponding tensor are + * derived from the NVTE_QKV_Layout enum (the historical behavior). `q`, `k` and `v` must + * either all be provided or all be NULL. `d_o` is only consulted by the backward pass. + * + * Only the F16 arbitrary-seqlen backend consumes these strides, and only for dense + * (non-THD, non-paged) layouts; the FP8 and max-512 backends, as well as THD/paged + * layouts, ignore them and always use the enum-derived strides. Output tensors + * (O, dQ/dK/dV) always keep the enum-derived strides since Transformer Engine + * allocates them. + */ +typedef struct NVTEQKVStrides { + const int64_t *q; /*!< Q strides [b, h, s, d] or NULL */ + const int64_t *k; /*!< K strides [b, h, s, d] or NULL */ + const int64_t *v; /*!< V strides [b, h, s, d] or NULL */ + const int64_t *d_o; /*!< dO strides [b, h, s, d] or NULL (backward only) */ +} NVTEQKVStrides; + +/*! \brief Compute dot product attention with separate Q, K and V, with explicit strides. + * + * \warning This API is **experimental** and subject to change. + * + * Identical to nvte_fused_attn_fwd(), with an additional `qkv_strides` parameter that + * provides the real memory strides of Q/K/V to the cuDNN graph instead of strides + * reconstructed from `qkv_layout` (see NVTEQKVStrides for the exact semantics). + * nvte_fused_attn_fwd() is equivalent to calling this function with all-NULL strides. + * + * \param[in] qkv_strides Real strides of Q/K/V; NULL members fall back + * to the enum-derived strides. + * + * All other parameters are as in nvte_fused_attn_fwd(). + */ +void nvte_fused_attn_fwd_v2( + 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, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream); + +/*! \brief Compute the backward of the dot product attention with separate Q, K and V, + * with explicit strides. + * + * \warning This API is **experimental** and subject to change. + * + * Identical to nvte_fused_attn_bwd(), with an additional `qkv_strides` parameter that + * provides the real memory strides of Q/K/V and dO to the cuDNN graph instead of strides + * reconstructed from `qkv_layout` (see NVTEQKVStrides for the exact semantics). + * nvte_fused_attn_bwd() is equivalent to calling this function with all-NULL strides. + * + * \param[in] qkv_strides Real strides of Q/K/V and dO; NULL members fall + * back to the enum-derived strides. + * + * All other parameters are as in nvte_fused_attn_bwd(). + */ +void nvte_fused_attn_bwd_v2( + 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, NVTEQKVStrides qkv_strides, + NVTETensor workspace, cudaStream_t stream); /*! \brief Get KV format for a given QKV layout. * diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 4763b58553..1d4b758fb2 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -14,9 +14,9 @@ namespace { constexpr int block_size = 512; -// Prototype (opt-in via NVTE_FUSED_ATTN_REAL_STRIDES=1): pass the real torch strides of -// Q/K/V (and dO in backward) down to the cuDNN fused-attention graph instead of relying -// on strides reconstructed from the NVTE_QKV_Layout enum. +// Opt-in via NVTE_FUSED_ATTN_REAL_STRIDES=1: pass the real torch strides of Q/K/V (and +// dO in backward) down to the cuDNN fused-attention graph via nvte_fused_attn_fwd/bwd_v2 +// instead of relying on strides reconstructed from the NVTE_QKV_Layout enum. bool real_strides_enabled() { const char *env = std::getenv("NVTE_FUSED_ATTN_REAL_STRIDES"); return env != nullptr && env[0] == '1'; @@ -293,33 +293,34 @@ std::vector fused_attn_fwd( // create workspace TensorWrapper workspace; - // Prototype: pass real torch strides of Q/K/V through to the cuDNN graph - bool real_strides_set = false; + // Real torch strides of Q/K/V for the cuDNN graph (opt-in); NULL members fall back to + // the enum-derived strides. + int64_t q_strides[4], k_strides[4], v_strides[4]; + NVTEQKVStrides qkv_strides{nullptr, nullptr, nullptr, nullptr}; if (real_strides_enabled()) { at::Tensor q_torch, k_torch, v_torch; if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && get_plain_torch_tensor(V, v_torch)) { NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - int64_t q_strides[4], k_strides[4], v_strides[4]; if (to_bhsd_strides(q_torch, q_format, q_strides) && to_bhsd_strides(k_torch, kv_format, k_strides) && to_bhsd_strides(v_torch, kv_format, v_strides)) { - nvte_fused_attn_set_strides(q_strides, k_strides, v_strides, nullptr); - real_strides_set = true; + qkv_strides = NVTEQKVStrides{q_strides, k_strides, v_strides, nullptr}; } } } // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( + nvte_fused_attn_fwd_v2( 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()); + window_size[1], bottom_right_diagonal, qkv_strides, workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace and auxiliary output tensors @@ -369,21 +370,17 @@ std::vector fused_attn_fwd( // execute the kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( + nvte_fused_attn_fwd_v2( 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()); + window_size[1], bottom_right_diagonal, qkv_strides, workspace.data(), + at::cuda::getCurrentCUDAStream()); }); - // clear the real-stride override - if (real_strides_set) { - nvte_fused_attn_set_strides(nullptr, nullptr, nullptr, nullptr); - } - // destroy tensor wrappers, but not allocated memory nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); @@ -643,36 +640,36 @@ std::vector fused_attn_bwd( // create workspace TensorWrapper workspace; - // Prototype: pass real torch strides of Q/K/V (and dO) through to the cuDNN graph - bool real_strides_set = false; + // Real torch strides of Q/K/V and dO for the cuDNN graph (opt-in); NULL members fall + // back to the enum-derived strides. + int64_t q_strides[4], k_strides[4], v_strides[4], do_strides[4]; + NVTEQKVStrides qkv_strides{nullptr, nullptr, nullptr, nullptr}; if (real_strides_enabled()) { at::Tensor q_torch, k_torch, v_torch, do_torch; if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && get_plain_torch_tensor(V, v_torch)) { - int64_t q_strides[4], k_strides[4], v_strides[4], do_strides[4]; if (to_bhsd_strides(q_torch, q_format, q_strides) && to_bhsd_strides(k_torch, kv_format, k_strides) && to_bhsd_strides(v_torch, kv_format, v_strides)) { const bool have_do_strides = get_plain_torch_tensor(dO, do_torch) && to_bhsd_strides(do_torch, do_format, do_strides); - nvte_fused_attn_set_strides(q_strides, k_strides, v_strides, - have_do_strides ? do_strides : nullptr); - real_strides_set = true; + qkv_strides = + NVTEQKVStrides{q_strides, k_strides, v_strides, have_do_strides ? do_strides : nullptr}; } } } // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( + nvte_fused_attn_bwd_v2( 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()); + window_size[1], bottom_right_diagonal, deterministic, cuda_graph, qkv_strides, + workspace.data(), at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace @@ -682,22 +679,17 @@ std::vector fused_attn_bwd( // execute kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( + nvte_fused_attn_bwd_v2( 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()); + window_size[1], bottom_right_diagonal, deterministic, cuda_graph, qkv_strides, + workspace.data(), at::cuda::getCurrentCUDAStream()); }); - // clear the real-stride override - if (real_strides_set) { - nvte_fused_attn_set_strides(nullptr, nullptr, nullptr, nullptr); - } - // destroy tensor wrappers nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); From aba881356fe9da4868004b19e9de5b8d7a86e1a5 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Tue, 7 Jul 2026 16:18:13 +0200 Subject: [PATCH 3/3] [PyTorch] fused attn: make real strides the default for dense F16, drop the opt-in flag Remove NVTE_FUSED_ATTN_REAL_STRIDES: the PyTorch extension now always passes the live torch strides of Q/K/V (and dO) to the cuDNN fused-attention graphs for plain 4-D torch tensors; quantized tensors, THD and non-4D inputs silently fall back to NULL strides (enum-derived). The v2 C API contract (NVTEQKVStrides, NULL = enum-derived) is unchanged. With real strides always on, DotProductAttention skips the pointer-based (untyped_storage().data_ptr() / storage_offset()) qkv layout detection for dense bshd/sbhd non-FP8 inputs and declares the format-derived separate layout instead. This removes the torch.compile graph breaks on UntypedStorage.data_ptr inside get_qkv_layout and the per-step CPU overhead of the detection. The old detection's stride normalization is preserved: tensors with stride(-1) != 1 are made contiguous in the bypass path (plain tensor metadata, still traceable). thd, FP8 DPA and KV-caching paths keep the existing detection since those backends do not consume real strides. Tests: DPA-level bit-exactness of packed-QKV strided views vs separate contiguous tensors (fused and flash backends), stride(-1) normalization, thd detection retention, and a torch.compile graph-break check with a forced-detection negative control. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Pawel Gadzinski --- .../attention/test_fused_attn_real_strides.py | 266 ++++++++++++++---- .../dot_product_attention.py | 45 +++ .../pytorch/csrc/extensions/attention.cpp | 22 +- 3 files changed, 265 insertions(+), 68 deletions(-) diff --git a/tests/pytorch/attention/test_fused_attn_real_strides.py b/tests/pytorch/attention/test_fused_attn_real_strides.py index f6589532b4..920463a3e7 100644 --- a/tests/pytorch/attention/test_fused_attn_real_strides.py +++ b/tests/pytorch/attention/test_fused_attn_real_strides.py @@ -2,23 +2,28 @@ # # See LICENSE for license information. -"""Tests for NVTE_FUSED_ATTN_REAL_STRIDES (nvte_fused_attn_fwd/bwd_v2). - -With the flag on, the PyTorch extension passes the real torch strides of -Q/K/V (and dO) to the cuDNN fused-attention graph instead of strides -reconstructed from the NVTE_QKV_Layout enum. Strided views into a packed -QKV buffer then compute correctly even when declared with the plain -*separate* layout enum -- the enum no longer needs to encode memory -geometry. +"""Tests for real-stride plumbing (nvte_fused_attn_fwd/bwd_v2). + +For dense (non-THD, non-paged) f16 layouts, the PyTorch extension passes +the real torch strides of Q/K/V (and dO) to the cuDNN fused-attention +graph instead of strides reconstructed from the NVTE_QKV_Layout enum. +Strided views into a packed QKV buffer then compute correctly even when +declared with the plain *separate* layout enum -- the enum no longer needs +to encode memory geometry -- and DotProductAttention no longer needs the +pointer-based (data_ptr/storage_offset) layout detection that graph-breaks +under torch.compile. """ -import os - import pytest import torch import transformer_engine.pytorch # noqa: F401 (loads libtransformer_engine.so) import transformer_engine_torch as tex +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, +) +import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils from transformer_engine.pytorch.cpp_extensions.fused_attn import ( fused_attn_bwd, fused_attn_fwd, @@ -31,18 +36,6 @@ _DTYPE = torch.bfloat16 -@pytest.fixture -def real_strides_flag(monkeypatch): - def set_flag(on: bool): - if on: - monkeypatch.setenv("NVTE_FUSED_ATTN_REAL_STRIDES", "1") - else: - monkeypatch.delenv("NVTE_FUSED_ATTN_REAL_STRIDES", raising=False) - - yield set_flag - monkeypatch.delenv("NVTE_FUSED_ATTN_REAL_STRIDES", raising=False) - - def _cu_seqlens(): return torch.arange(0, (_B + 1) * _S, _S, dtype=torch.int32, device="cuda") @@ -85,38 +78,19 @@ def _backend_supported(): @requires_backend -@pytest.mark.parametrize("fmt", ["bshd", "sbhd"]) -def test_contiguous_separate_matches_enum_path(real_strides_flag, fmt): - """Contiguous separate q/k/v: real strides == enum strides, so flag on == flag off.""" - torch.manual_seed(0) - shape = (_B, _S, _H, _D) if fmt == "bshd" else (_S, _B, _H, _D) - q = torch.randn(*shape, dtype=_DTYPE, device="cuda") - k, v, d_o = q.clone(), q.clone(), torch.randn(*shape, dtype=_DTYPE, device="cuda") - layout = f"{fmt}_{fmt}_{fmt}" - - real_strides_flag(False) - reference = _run(q, k, v, d_o, layout, fmt) - real_strides_flag(True) - result = _run(q, k, v, d_o, layout, fmt) - _assert_bit_exact(result, reference) - - -@requires_backend -def test_packed_bs3hd_views_declared_separate(real_strides_flag): +def test_packed_bs3hd_views_declared_separate(): """Headline: strided views into a packed [b,s,3,h,d] buffer, declared with the - plain separate layout enum, are bit-exact vs the contiguous baseline when real - strides are passed.""" + plain separate layout enum, are bit-exact vs the contiguous baseline because the + real strides are passed to the cuDNN graph.""" torch.manual_seed(0) qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") d_o = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") - real_strides_flag(False) reference = _run( qkv[:, :, 0].contiguous(), qkv[:, :, 1].contiguous(), qkv[:, :, 2].contiguous(), d_o, "bshd_bshd_bshd", "bshd", ) - real_strides_flag(True) q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] assert not q.is_contiguous() result = _run(q, k, v, d_o, "bshd_bshd_bshd", "bshd") @@ -124,19 +98,17 @@ def test_packed_bs3hd_views_declared_separate(real_strides_flag): @requires_backend -def test_packed_sbh3d_views_declared_separate(real_strides_flag): +def test_packed_sbh3d_views_declared_separate(): """Same as above for the sbh3d interleave (packing at dim -2, sbhd format).""" torch.manual_seed(0) qkv = torch.randn(_S, _B, _H, 3, _D, dtype=_DTYPE, device="cuda") d_o = torch.randn(_S, _B, _H, _D, dtype=_DTYPE, device="cuda") - real_strides_flag(False) reference = _run( qkv[:, :, :, 0].contiguous(), qkv[:, :, :, 1].contiguous(), qkv[:, :, :, 2].contiguous(), d_o, "sbhd_sbhd_sbhd", "sbhd", ) - real_strides_flag(True) q, k, v = qkv[:, :, :, 0], qkv[:, :, :, 1], qkv[:, :, :, 2] assert not q.is_contiguous() result = _run(q, k, v, d_o, "sbhd_sbhd_sbhd", "sbhd") @@ -144,16 +116,206 @@ def test_packed_sbh3d_views_declared_separate(real_strides_flag): @requires_backend -def test_packed_views_with_packed_enum_unchanged(real_strides_flag): +def test_packed_views_with_packed_enum_unchanged(): """Regression: the historical path (packed views with the matching packed enum) - is unchanged with the flag on.""" + still computes the same values as the separate-contiguous baseline. For packed + views the real strides coincide with the enum-derived ones, so passing them is + a no-op numerically.""" torch.manual_seed(0) qkv = torch.randn(_B, _S, 3, _H, _D, dtype=_DTYPE, device="cuda") d_o = torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda") q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] - real_strides_flag(False) - reference = _run(q, k, v, d_o, "bs3hd", "bshd") - real_strides_flag(True) + reference = _run( + qkv[:, :, 0].contiguous(), qkv[:, :, 1].contiguous(), qkv[:, :, 2].contiguous(), + d_o, "bshd_bshd_bshd", "bshd", + ) result = _run(q, k, v, d_o, "bs3hd", "bshd") _assert_bit_exact(result, reference) + + +# --------------------------------------------------------------------------- +# DotProductAttention end-to-end: DPA skips pointer-based qkv layout detection +# (data_ptr/storage_offset games) for dense f16 layouts and declares the +# format-derived separate layout instead. +# --------------------------------------------------------------------------- + + +def _force_backend(monkeypatch, backend): + """Force a single attention backend via env and invalidate the selection cache.""" + flash, fused = {"flash": ("1", "0"), "fused": ("0", "1")}[backend] + monkeypatch.setenv("NVTE_FLASH_ATTN", flash) + monkeypatch.setenv("NVTE_FUSED_ATTN", fused) + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + if backend == "flash": + # flash-attn bwd uses atomics unless deterministic; the fused (cuDNN) + # backend gets disabled outright by the deterministic flag on some + # devices, and its bwd is run-to-run deterministic anyway. + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + dpa_module._attention_backends["backend_selection_requires_update"] = True + + +def _make_dpa(qkv_format): + return DotProductAttention( + _H, _D, attention_dropout=0.0, qkv_format=qkv_format, attn_mask_type="no_mask" + ) + + +def _dpa_separate(qkv_format): + """Fwd+bwd on three contiguous leaves; returns (out, dq, dk, dv).""" + torch.manual_seed(0) + shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + q, k, v = [ + torch.randn(*shape, dtype=_DTYPE, device="cuda", requires_grad=True) for _ in range(3) + ] + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(q, k, v) + out.backward(torch.ones_like(out)) + return out, q.grad, k.grad, v.grad + + +def _dpa_packed_views(qkv_format): + """Fwd+bwd on strided views into one packed leaf; returns (out, dq, dk, dv).""" + torch.manual_seed(0) + shape = (_B, _S, _H, _D) if qkv_format == "bshd" else (_S, _B, _H, _D) + parts = [torch.randn(*shape, dtype=_DTYPE, device="cuda") for _ in range(3)] + qkv = torch.stack(parts, dim=2).requires_grad_() # bs3hd / sb3hd packing + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + assert not q.is_contiguous() + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa(qkv_format)(q, k, v) + out.backward(torch.ones_like(out)) + return out, qkv.grad[:, :, 0], qkv.grad[:, :, 1], qkv.grad[:, :, 2] + + +@requires_backend +@pytest.mark.parametrize("fmt", ["bshd", "sbhd"]) +def test_dpa_fused_packed_views(monkeypatch, fmt): + """Fused backend: packed stack views (declared separate, real strides) are + bit-exact vs contiguous separate q/k/v, fwd and grads.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate(fmt) + result = _dpa_packed_views(fmt) + _assert_bit_exact(result, reference) + + +def test_dpa_flash_packed_views(monkeypatch): + """Flash backend smoke: flash consumes real strides natively, so packed views + declared separate are bit-exact vs separate.""" + _force_backend(monkeypatch, "flash") + try: + reference = _dpa_separate("bshd") + except Exception as exc: + pytest.skip(f"flash attention backend not available: {exc}") + result = _dpa_packed_views("bshd") + _assert_bit_exact(result, reference) + + +@requires_backend +def test_dpa_layout_detection_skipped_dense_kept_thd(monkeypatch): + """get_qkv_layout is not called for dense bshd, but still runs for thd + (ragged layouts ignore strides in C++, so detection stays).""" + _force_backend(monkeypatch, "fused") + + calls = [] + orig = dpa_utils.get_qkv_layout + + def counting(*args, **kwargs): + calls.append(kwargs.get("qkv_format")) + return orig(*args, **kwargs) + + monkeypatch.setattr(dpa_utils, "get_qkv_layout", counting) + + _dpa_separate("bshd") + assert not calls, "get_qkv_layout should be skipped for dense bshd" + + # thd: full sequences, padding mask + torch.manual_seed(0) + t = _B * _S + q, k, v = [torch.randn(t, _H, _D, dtype=_DTYPE, device="cuda") for _ in range(3)] + cu = _cu_seqlens() + dpa_module._attention_backends["backend_selection_requires_update"] = True + dpa = DotProductAttention( + _H, _D, attention_dropout=0.0, qkv_format="thd", attn_mask_type="padding" + ) + try: + dpa(q, k, v, cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=_S, max_seqlen_kv=_S) + except ValueError: + # No thd-capable backend on this device; the layout step (the subject + # of this test) runs before backend dispatch, so the assertion below + # still holds. + pass + assert calls == ["thd"], "get_qkv_layout must still run for thd" + + +@requires_backend +def test_dpa_noncontiguous_head_dim_normalized(monkeypatch): + """Inputs with stride(-1) != 1 are normalized with .contiguous() in the bypass + path (mirrors the old detection's contiguous-retry) and stay bit-exact.""" + _force_backend(monkeypatch, "fused") + reference = _dpa_separate("bshd") + + torch.manual_seed(0) + q, k, v = [ + torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) + for _ in range(3) + ] + # transpose(-1, -2) of a transposed copy: same values, stride(-1) != 1 + q_t = q.detach().transpose(2, 3).contiguous().transpose(2, 3).requires_grad_() + k_t = k.detach().transpose(2, 3).contiguous().transpose(2, 3).requires_grad_() + v_t = v.detach().transpose(2, 3).contiguous().transpose(2, 3).requires_grad_() + assert q_t.stride(-1) != 1 + dpa_module._attention_backends["backend_selection_requires_update"] = True + out = _make_dpa("bshd")(q_t, k_t, v_t) + out.backward(torch.ones_like(out)) + _assert_bit_exact((out, q_t.grad, k_t.grad, v_t.grad), reference) + + +def _compiled_graph_breaks(): + """Compile DPA (after an eager warm-up that caches backend selection) and + return dynamo's graph_break counters for one fwd+bwd.""" + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + torch.manual_seed(0) + q, k, v = [ + torch.randn(_B, _S, _H, _D, dtype=_DTYPE, device="cuda", requires_grad=True) + for _ in range(3) + ] + dpa = _make_dpa("bshd") + dpa_module._attention_backends["backend_selection_requires_update"] = True + dpa(q, k, v) # eager warm-up: backend selection happens outside dynamo + out = torch.compile(dpa)(q, k, v) + out.backward(torch.ones_like(out)) + breaks = dict(torch._dynamo.utils.counters["graph_break"]) + torch._dynamo.reset() + return breaks + + +def _pointer_breaks(breaks): + return { + reason: count + for reason, count in breaks.items() + if "data_ptr" in reason or "UntypedStorage" in reason + } + + +@requires_backend +def test_dpa_torch_compile_no_data_ptr_graph_breaks(monkeypatch): + """Compiling DPA no longer graph-breaks on data_ptr/UntypedStorage: the + pointer-based layout detection is bypassed for dense layouts. Negative + control: forcing the detection back on reintroduces those breaks.""" + _force_backend(monkeypatch, "fused") + + # negative control: force pointer-based detection, expect data_ptr breaks + monkeypatch.setattr(dpa_module, "_skip_pointer_layout_detection", lambda *a: False) + baseline = _compiled_graph_breaks() + assert _pointer_breaks(baseline), ( + f"expected data_ptr/UntypedStorage breaks with detection forced on: {baseline}" + ) + monkeypatch.undo() + + _force_backend(monkeypatch, "fused") + breaks = _compiled_graph_breaks() + assert not _pointer_breaks(breaks), ( + f"data_ptr/UntypedStorage graph breaks should be gone: {_pointer_breaks(breaks)}" + ) 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 03008bb2d7..7ae2610947 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 @@ -190,6 +190,33 @@ def _pad_qkv_head_dim(query_layer, key_layer, value_layer): return query_layer, key_layer, value_layer, orig_head_dim_qk, orig_head_dim_v +def _skip_pointer_layout_detection(qkv_format, fp8_dpa, inference_params) -> bool: + """Decide whether pointer-based qkv layout detection can be skipped. + + ``get_qkv_layout`` inspects ``untyped_storage().data_ptr()`` and + ``storage_offset()`` of q/k/v on every forward to detect packed layouts + (e.g. ``bs3hd``). This is torch.compile-hostile (data_ptr access causes + graph breaks) and adds per-step CPU overhead. For dense f16 layouts the + detection is redundant for memory correctness: the cuDNN fused-attention + backend receives the live tensor strides via the v2 C API, + flash-attention consumes real strides natively, and the unfused backend + uses plain torch ops. The format-derived separate layout string is then + always safe to declare. + + Excluded cases (detection must still run): + - ``thd``: the C++ side ignores strides for ragged layouts, so a packed + ``t3hd``/``th3d`` input declared as ``thd_thd_thd`` would silently + compute on wrong memory. + - FP8 DPA: the FP8 fused backend does not consume the real-strides + parameter, and ``combine_and_quantize`` relies on detected packedness. + - KV caching (``inference_params``): layouts may need the ``paged_kv_`` + prefix and mixed q/kv formats that only detection derives. + """ + if inference_params is not None or fp8_dpa: + return False + return qkv_format in ("bshd", "sbhd") + + def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim_v): """Trim FlashAttention output after padding V to a larger head dimension.""" out_shape = attn_out.shape[:-1] @@ -1449,6 +1476,24 @@ def forward( qkv_format=qkv_format, inference_params=inference_params, ) + elif _skip_pointer_layout_detection( + qkv_format, + self.fp8 and self.fp8_meta["recipe"].fp8_dpa, + inference_params, + ): + # The backends consume the live tensor strides for dense + # layouts, so the layout enum no longer needs to encode the + # memory geometry: declare the format-derived separate layout + # and pass q/k/v through as-is. Only the last (head) dimension + # must be packed; normalize it if needed (stride(-1) is plain + # tensor metadata, so this stays torch.compile-traceable). + query_layer, key_layer, value_layer = [ + x if x.stride(-1) == 1 else x.contiguous() + for x in [query_layer, key_layer, value_layer] + ] + qkv_layout = f"{qkv_format}_{qkv_format}_{qkv_format}" + q_format = qkv_format + kv_format = qkv_format else: ( qkv_layout, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 1d4b758fb2..310eb8a8b6 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -4,8 +4,6 @@ * See LICENSE for license information. ************************************************************************/ -#include - #include "../extensions.h" #include "common.h" #include "pybind.h" @@ -14,14 +12,6 @@ namespace { constexpr int block_size = 512; -// Opt-in via NVTE_FUSED_ATTN_REAL_STRIDES=1: pass the real torch strides of Q/K/V (and -// dO in backward) down to the cuDNN fused-attention graph via nvte_fused_attn_fwd/bwd_v2 -// instead of relying on strides reconstructed from the NVTE_QKV_Layout enum. -bool real_strides_enabled() { - const char *env = std::getenv("NVTE_FUSED_ATTN_REAL_STRIDES"); - return env != nullptr && env[0] == '1'; -} - // Cast a py::handle to at::Tensor if it wraps a plain torch tensor. bool get_plain_torch_tensor(pybind11::handle handle, at::Tensor &out) { try { @@ -293,11 +283,11 @@ std::vector fused_attn_fwd( // create workspace TensorWrapper workspace; - // Real torch strides of Q/K/V for the cuDNN graph (opt-in); NULL members fall back to - // the enum-derived strides. + // Real torch strides of Q/K/V for the cuDNN graph; NULL members fall back to + // the enum-derived strides (quantized tensors, THD, non-4D tensors). int64_t q_strides[4], k_strides[4], v_strides[4]; NVTEQKVStrides qkv_strides{nullptr, nullptr, nullptr, nullptr}; - if (real_strides_enabled()) { + { at::Tensor q_torch, k_torch, v_torch; if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && get_plain_torch_tensor(V, v_torch)) { @@ -640,11 +630,11 @@ std::vector fused_attn_bwd( // create workspace TensorWrapper workspace; - // Real torch strides of Q/K/V and dO for the cuDNN graph (opt-in); NULL members fall - // back to the enum-derived strides. + // Real torch strides of Q/K/V and dO for the cuDNN graph; NULL members fall + // back to the enum-derived strides (quantized tensors, THD, non-4D tensors). int64_t q_strides[4], k_strides[4], v_strides[4], do_strides[4]; NVTEQKVStrides qkv_strides{nullptr, nullptr, nullptr, nullptr}; - if (real_strides_enabled()) { + { at::Tensor q_torch, k_torch, v_torch, do_torch; if (get_plain_torch_tensor(Q, q_torch) && get_plain_torch_tensor(K, k_torch) && get_plain_torch_tensor(V, v_torch)) {