From 5e93df9f92009a26d370c0c7f6150be81136a331 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:10:08 +0000 Subject: [PATCH 1/3] Add packed attention workspace recipes Introduce checked, graph-free workspace recipes for PackedAttention and PackedMultiHeadAttention while preserving legacy allocation sizes and backend layouts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78 --- cmake/onnxruntime_unittests.cmake | 1 + .../contrib_ops/cuda/bert/attention_data.h | 39 - .../bert/cutlass_fmha/fmha_launch_template.h | 20 +- .../contrib_ops/cuda/bert/packed_attention.cc | 211 +-- .../contrib_ops/cuda/bert/packed_attention.h | 17 +- .../cuda/bert/packed_attention_data.h | 58 + .../cuda/bert/packed_attention_impl.cu | 101 +- .../cuda/bert/packed_attention_impl.h | 25 +- .../cuda/bert/packed_attention_workspace.cc | 1309 +++++++++++++++++ .../cuda/bert/packed_attention_workspace.h | 226 +++ .../cuda/bert/packed_multihead_attention.cc | 181 +-- .../cuda/bert/packed_multihead_attention.h | 9 +- .../bert/packed_multihead_attention_impl.cu | 89 +- .../bert/packed_multihead_attention_impl.h | 2 +- .../contrib_ops/packed_attention_op_test.cc | 163 ++ .../packed_multihead_attention_op_test.cc | 243 ++- .../packed_attention_workspace_header_test.cc | 40 + .../packed_attention_workspace_test.cc | 914 ++++++++++++ 18 files changed, 3162 insertions(+), 486 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention_data.h create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.cc create mode 100644 onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.h create mode 100644 onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_header_test.cc create mode 100644 onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_test.cc diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index a64661c58aee5..a9141f116a8f0 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1042,6 +1042,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P set(onnxruntime_test_providers_cuda_plugin_internal_test_src "${TEST_SRC_DIR}/providers/cuda/test_cases/allocator_cuda_test.cc" "${TEST_SRC_DIR}/providers/cuda/test_cases/cuda_utils_test.cc" + "${TEST_SRC_DIR}/providers/cuda/test_cases/packed_attention_workspace_header_test.cc" "${TEST_SRC_DIR}/providers/cuda/test_cases/reduction_functions_test.cc" ) # matmul_nbits_workspace_test.cc / matmul_nbits_e2e_workspace_test.cc are intentionally excluded diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 6d8448b0eb65f..9cb66f383bd36 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -106,45 +106,6 @@ struct AttentionData { } }; -template -struct PackedAttentionData { - T* gemm_buffer; - const T* bias; - const T* attention_bias; - const int32_t* token_offset; - const int32_t* cumulative_sequence_length; - - T* workspace; - T* output; - - void* fused_runner; - - bool use_memory_efficient_attention; -}; - -template -struct PackedMultiHeadAttentionData { - const T* query; - const T* key; - const T* value; - const T* bias; - const T* attention_bias; - - const int32_t* token_offset; - const int32_t* cumulative_sequence_length; - - AttentionQkvFormat source_qkv_format; - - bool no_qkv_workspace; - T* workspace; - T* output; - - void* fused_runner; - - bool use_flash_attention; - bool use_memory_efficient_attention; -}; - template struct GroupQueryAttentionData { // Input Tensors diff --git a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h index 41bf5e1bbfb0d..6eafc9dc7ccf4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h +++ b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h @@ -220,11 +220,15 @@ void LaunchCutlassFmha(const MemoryEfficientAttentionParams& params) { } if (params.attn_bias != nullptr) { - p.bias_strideH = params.broadcast_attn_bias_dim_1 ? 0 : p.num_queries * p.num_keys; + p.bias_strideH = params.broadcast_attn_bias_dim_1 + ? 0 + : static_cast(p.num_queries) * p.num_keys; p.bias_strideM = p.num_keys; p.bias_strideB = params.broadcast_attn_bias_dim_0 ? 0 - : ((params.broadcast_attn_bias_dim_1 ? 1 : params.num_heads) * p.num_queries * p.num_keys); + : static_cast( + params.broadcast_attn_bias_dim_1 ? 1 : params.num_heads) * + p.num_queries * p.num_keys; } else { p.bias_strideH = 0; p.bias_strideM = 0; @@ -297,10 +301,14 @@ void DispatchIsAligned(const MemoryEfficientAttentionParams& params) { int num_queries = params.sequence_length; int bias_strideM = num_keys; // Broadcast dimensions use stride=0, which satisfies any alignment (0 % N == 0). - int bias_strideH = params.broadcast_attn_bias_dim_1 ? 0 : num_queries * num_keys; - int bias_strideB = params.broadcast_attn_bias_dim_0 - ? 0 - : ((params.broadcast_attn_bias_dim_1 ? 1 : params.num_heads) * num_queries * num_keys); + int64_t bias_strideH = params.broadcast_attn_bias_dim_1 + ? 0 + : static_cast(num_queries) * num_keys; + int64_t bias_strideB = params.broadcast_attn_bias_dim_0 + ? 0 + : static_cast( + params.broadcast_attn_bias_dim_1 ? 1 : params.num_heads) * + num_queries * num_keys; is_aligned = is_aligned && bias_strideM % AlignedAK::kAlignmentQ == 0 && (params.num_heads <= 1 || bias_strideH % AlignedAK::kAlignmentQ == 0) && diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc index 68140390bd132..1508d646aa3f0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc @@ -9,7 +9,6 @@ #include "contrib_ops/cuda/bert/packed_attention_impl.h" #include "contrib_ops/cuda/bert/bert_padding.h" #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" -#include "contrib_ops/cpu/bert/multihead_attention_helper.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -19,6 +18,27 @@ namespace onnxruntime { namespace contrib { namespace cuda { +PackedAttentionShape MakePackedAttentionShape(const TensorShape& shape) noexcept { + PackedAttentionShape result; + const auto& dimensions = shape.GetDims(); + result.rank = dimensions.size(); + const size_t dimensions_to_copy = + dimensions.size() < result.dimensions.size() ? dimensions.size() : result.dimensions.size(); + for (size_t i = 0; i < dimensions_to_copy; ++i) { + result.dimensions[i] = dimensions[i]; + } + + return result; +} + +Status PackedAttentionWorkspaceStatusToStatus(PackedAttentionWorkspaceStatus status) { + if (status.IsOK()) { + return Status::OK(); + } + + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, status.message); +} + #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ PackedAttention, \ @@ -90,7 +110,7 @@ PackedAttention::PackedAttention(const OpKernelInfo& info) : TrtFusedAttention(info) { int64_t num_heads = 0; ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); - num_heads_ = static_cast(num_heads); + num_heads_ = num_heads; scale_ = info.GetAttrOrDefault("scale", 0.0f); @@ -106,124 +126,42 @@ Status PackedAttention::CheckInputs(const TensorShape& input_shape, const TensorShape& token_offset_shape, const TensorShape& cu_seq_len_shape, const Tensor* attention_bias, - PackedAttentionParameters& parameters) const { - // Abbreviation and Meanings: - // T: token_count - // B: batch_size - // S: sequence_length - // N: num_heads - // H: head size for Q and K, aka q_head_size or v_head_size or qk_head_size - // H_v: v_head_size - // D_i: input hidden size - // D: hidden size for Q and K (D = N * H), aka q_hidden_size or k_hidden_size or qk_hidden_size - // D_v: v_hidden_size = num_heads * v_head_size - - // Input shapes: - // input: : (T, D_i) - // weights (Q/K/V) : (D_i, D + D + D_v) - // bias (Q/K/V) : (D + D + D_v) - // token_offset : (B, S) - // cu_seq_len_shape : (B + 1) - // attention_bias : (B or 1, N or 1, S, S) or NULL - const auto& input_dims = input_shape.GetDims(); - if (input_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'input' is expected to have 2 dimensions in packing mode, got ", - input_dims.size()); - } - int64_t token_count = input_dims[0]; - int64_t input_hidden_size = input_dims[1]; - - const auto& token_offset_dims = token_offset_shape.GetDims(); - if (token_offset_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'packing_token_offset' is expected to have 2 dimensions in packing mode, got ", - token_offset_dims.size()); - } - - int64_t batch_size = token_offset_dims[0]; - int64_t sequence_length = token_offset_dims[1]; - - const auto& bias_dims = bias_shape.GetDims(); - if (bias_dims.size() != 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' is expected to have 1 dimension, got ", - bias_dims.size()); - } - - const auto& weights_dims = weights_shape.GetDims(); - if (weights_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'weights' is expected to have 2 dimensions, got ", - weights_dims.size()); - } - if (weights_dims[0] != input_hidden_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 1 dimension 0 should have same length as dimension 2 of input 0"); - } - - if (bias_dims[0] != weights_dims[1]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'bias' dimension 0 should have same length as dimension 1 of input 'weights'"); - } - - const auto& cu_seq_len_dims = cu_seq_len_shape.GetDims(); - if (cu_seq_len_dims.size() != 1 || cu_seq_len_dims[0] != batch_size + 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'cumulative_sequence_length' should have 1 dimension with size equal to batch_size + 1"); - } - - const int num_heads = this->GetNumHeads(); - int64_t q_hidden_size = bias_dims[0] / static_cast(3); - int64_t k_hidden_size = q_hidden_size; - int64_t v_hidden_size = k_hidden_size; - if (qkv_hidden_sizes_.size() != 0) { - if (qkv_hidden_sizes_.size() != 3) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "qkv_hidden_sizes attribute should have 3 elements"); - } - - for (size_t i = 0; i < qkv_hidden_sizes_.size(); i++) { - if (qkv_hidden_sizes_[i] % num_heads != 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "hidden_size should be divisible by num_heads:", qkv_hidden_sizes_[i]); - } - } - - q_hidden_size = qkv_hidden_sizes_[0]; - k_hidden_size = qkv_hidden_sizes_[1]; - v_hidden_size = qkv_hidden_sizes_[2]; - } - - if (q_hidden_size != k_hidden_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "qkv_hidden_sizes first element should be same as the second"); - } - - if (bias_dims[0] != q_hidden_size + k_hidden_size + v_hidden_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'bias' dimension 0 should have same length as sum of Q/K/V hidden sizes:", - " q_hidden_size=", q_hidden_size, " k_hidden_size=", k_hidden_size, " v_hidden_size=", - v_hidden_size, "bias_dims[0]=", bias_dims[0]); + PackedAttentionParameters& parameters, + PackedAttentionProblem& problem) const { + PackedAttentionInputShapes inputs; + inputs.input = MakePackedAttentionShape(input_shape); + inputs.weights = MakePackedAttentionShape(weights_shape); + inputs.bias = MakePackedAttentionShape(bias_shape); + inputs.token_offset = MakePackedAttentionShape(token_offset_shape); + inputs.cumulative_sequence_length = MakePackedAttentionShape(cu_seq_len_shape); + inputs.element_size = sizeof(T); + inputs.num_heads = GetNumHeads(); + inputs.qkv_hidden_sizes_count = qkv_hidden_sizes_.size(); + for (size_t i = 0; i < qkv_hidden_sizes_.size() && i < inputs.qkv_hidden_sizes.size(); ++i) { + inputs.qkv_hidden_sizes[i] = qkv_hidden_sizes_[i]; } - - gsl::span attention_bias_dims; + inputs.has_attention_bias = attention_bias != nullptr; if (attention_bias != nullptr) { - attention_bias_dims = attention_bias->Shape().GetDims(); - ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckAttentionBias( - attention_bias_dims, batch_size, num_heads, sequence_length, sequence_length)); + inputs.attention_bias = MakePackedAttentionShape(attention_bias->Shape()); } - parameters.broadcast_attn_bias_dim_0 = attention_bias_dims.size() > 0 && attention_bias_dims[0] == 1; - parameters.broadcast_attn_bias_dim_1 = attention_bias_dims.size() > 1 && attention_bias_dims[1] == 1; - - parameters.batch_size = static_cast(batch_size); - parameters.sequence_length = static_cast(sequence_length); - parameters.input_hidden_size = static_cast(input_hidden_size); - parameters.hidden_size = static_cast(q_hidden_size); - parameters.v_hidden_size = static_cast(v_hidden_size); - parameters.head_size = static_cast(q_hidden_size) / num_heads; - parameters.v_head_size = static_cast(v_hidden_size) / num_heads; - parameters.num_heads = num_heads; + + auto problem_result = BuildPackedAttentionProblem(inputs); + ORT_RETURN_IF_ERROR(PackedAttentionWorkspaceStatusToStatus(problem_result.status)); + problem = problem_result.problem; + + parameters.broadcast_attn_bias_dim_0 = problem.broadcast_attn_bias_dim_0; + parameters.broadcast_attn_bias_dim_1 = problem.broadcast_attn_bias_dim_1; + + parameters.batch_size = problem.batch_size; + parameters.sequence_length = problem.sequence_length; + parameters.input_hidden_size = problem.input_hidden_size; + parameters.hidden_size = problem.hidden_size; + parameters.v_hidden_size = problem.v_hidden_size; + parameters.head_size = problem.qk_head_size; + parameters.v_head_size = problem.v_head_size; + parameters.num_heads = problem.num_heads; parameters.scale = this->GetScale(); - parameters.token_count = static_cast(token_count); + parameters.token_count = problem.token_count; return Status::OK(); } @@ -238,6 +176,7 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { const Tensor* attention_bias = context->Input(5); PackedAttentionParameters parameters; + PackedAttentionProblem problem; parameters.use_tf32 = this->UseTF32(); ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), weights->Shape(), @@ -245,7 +184,8 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { token_offset->Shape(), cumulative_sequence_length->Shape(), attention_bias, - parameters)); + parameters, + problem)); TensorShapeVector output_shape{parameters.token_count, parameters.v_hidden_size}; Tensor* output = context->Output(0, output_shape); @@ -281,11 +221,21 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { CudaT one = ToCudaType::FromFloat(1.0f); CudaT zero = ToCudaType::FromFloat(0.0f); - IAllocatorUniquePtr gemm_buffer; - int m = parameters.token_count; - int n = parameters.hidden_size + parameters.hidden_size + parameters.v_hidden_size; - int k = parameters.input_hidden_size; - gemm_buffer = this->template GetScratchBuffer(static_cast(m) * n, this->GetComputeStream(context)); + problem.backend = fused_runner != nullptr + ? PackedAttentionBackend::Trt + : (use_memory_efficient_attention + ? PackedAttentionBackend::MemoryEfficient + : PackedAttentionBackend::Unfused); + problem.trt_runner_available = fused_runner != nullptr; + auto workspace_result = GetPackedAttentionWorkspaceRecipe(problem); + ORT_RETURN_IF_ERROR(PackedAttentionWorkspaceStatusToStatus(workspace_result.status)); + const PackedAttentionWorkspaceRecipe& workspace_recipe = workspace_result.recipe; + + auto gemm_buffer = this->template GetScratchBuffer( + workspace_recipe.projection_bytes, this->GetComputeStream(context)); + const int m = workspace_recipe.projection_m; + const int n = workspace_recipe.projection_n; + const int k = workspace_recipe.projection_k; cublasHandle_t cublas = this->GetCublasHandle(context); @@ -297,21 +247,9 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { reinterpret_cast(input->Data()), k, &zero, reinterpret_cast(gemm_buffer.get()), n, device_prop, this->UseTF32())); - constexpr size_t element_size = sizeof(T); - constexpr bool no_qkv_workspace = false; // need workspace to add bias - size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, - parameters.batch_size, - parameters.num_heads, - parameters.head_size, - parameters.v_head_size, - parameters.sequence_length, - fused_runner, - false, - use_memory_efficient_attention, - no_qkv_workspace); - auto work_space = this->template GetScratchBuffer(workSpaceSize, this->GetComputeStream(context)); + auto work_space = this->template GetScratchBuffer( + workspace_recipe.attention_workspace_bytes, this->GetComputeStream(context)); - typedef typename ToCudaType::MappedType CudaT; PackedAttentionData data; data.gemm_buffer = reinterpret_cast(gemm_buffer.get()); data.bias = reinterpret_cast(bias->Data()); @@ -322,6 +260,7 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { data.output = reinterpret_cast(output->MutableData()); data.fused_runner = reinterpret_cast(fused_runner); data.use_memory_efficient_attention = use_memory_efficient_attention; + data.workspace_recipe = workspace_recipe; return QkvToContext(device_prop, cublas, this->Stream(context), parameters, data); } diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention.h index 0f7456d0cd0d8..c134009d198ba 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention.h @@ -5,12 +5,12 @@ #include #include - #include "core/providers/cuda/cuda_kernel.h" #include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h" #include "contrib_ops/cpu/bert/attention_common.h" #include "contrib_ops/cpu/bert/attention_parameters.h" #include "contrib_ops/cuda/bert/attention_kernel_options.h" +#include "contrib_ops/cuda/bert/packed_attention_workspace.h" namespace onnxruntime { namespace contrib { @@ -18,6 +18,10 @@ namespace cuda { using namespace onnxruntime::cuda; +PackedAttentionShape MakePackedAttentionShape(const TensorShape& shape) noexcept; + +Status PackedAttentionWorkspaceStatusToStatus(PackedAttentionWorkspaceStatus status); + template class TrtFusedAttention : public CudaKernel { public: @@ -50,15 +54,16 @@ class PackedAttention final : public TrtFusedAttention { const TensorShape& packing_token_offset_shape, const TensorShape& cu_seq_len_shape, const Tensor* attention_bias, - PackedAttentionParameters& parameters) const; + PackedAttentionParameters& parameters, + PackedAttentionProblem& problem) const; - int GetNumHeads() const { return num_heads_; } + int64_t GetNumHeads() const { return num_heads_; } float GetScale() const { return scale_; } private: - int num_heads_; // number of attention heads - float scale_; // scale for softmax. Default is 0.0f, which will be replaced by 1/sqrt(num_heads) later - std::vector qkv_hidden_sizes_; // Q, K, V hidden sizes parsed from the qkv_hidden_sizes attribute. + int64_t num_heads_; // number of attention heads + float scale_; // scale for softmax + std::vector qkv_hidden_sizes_; // Q, K, V hidden sizes }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_data.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention_data.h new file mode 100644 index 0000000000000..697506148839f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_data.h @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "contrib_ops/cpu/bert/attention_common.h" +#include "contrib_ops/cuda/bert/packed_attention_workspace.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +struct PackedAttentionData { + T* gemm_buffer; + const T* bias; + const T* attention_bias; + const int32_t* token_offset; + const int32_t* cumulative_sequence_length; + + T* workspace; + T* output; + + void* fused_runner; + + bool use_memory_efficient_attention; + PackedAttentionWorkspaceRecipe workspace_recipe; +}; + +template +struct PackedMultiHeadAttentionData { + const T* query; + const T* key; + const T* value; + const T* bias; + const T* attention_bias; + + const int32_t* token_offset; + const int32_t* cumulative_sequence_length; + + AttentionQkvFormat source_qkv_format; + + bool no_qkv_workspace; + T* workspace; + T* output; + + void* fused_runner; + + bool use_flash_attention; + bool use_memory_efficient_attention; + PackedAttentionWorkspaceRecipe workspace_recipe; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu index 7115e1da64715..8bfd778fcc061 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu @@ -26,62 +26,8 @@ namespace onnxruntime { namespace contrib { namespace cuda { -constexpr size_t kCUDAMemoryAlignment = 256; - constexpr int32_t kMAX_THREADS_PER_BLOCK = 256; -size_t GetAttentionScratchSize( - size_t element_size, - size_t batch_size, - size_t num_heads, - size_t sequence_length) { - const size_t bytes = element_size * batch_size * num_heads * sequence_length * sequence_length; - return ((bytes + kCUDAMemoryAlignment - 1) / kCUDAMemoryAlignment) * kCUDAMemoryAlignment; -} - -size_t GetAttentionWorkspaceSize( - size_t element_size, - size_t batch_size, - size_t num_heads, - size_t qk_head_size, - size_t v_head_size, - size_t sequence_length, - void* fused_runner, - bool use_flash_attention, - bool use_memory_efficient_attention, - bool no_qkv_workspace) { - // Note that q, k and v might need alignment for fused attention kernels. - const size_t qkv_bytes = no_qkv_workspace ? 0 : (element_size * batch_size * num_heads * sequence_length * (qk_head_size + qk_head_size + v_head_size)); - -#if USE_FLASH_ATTENTION - // Use portion of workspace for softmax buffer. - if (use_flash_attention) { - size_t flash_buffer_bytes = onnxruntime::flash::get_softmax_lse_size(sequence_length, batch_size, num_heads); - return qkv_bytes + flash_buffer_bytes; - } -#else - ORT_UNUSED_PARAMETER(use_flash_attention); -#endif - - if (fused_runner != nullptr) { - return qkv_bytes; - } - -#if USE_MEMORY_EFFICIENT_ATTENTION - if (use_memory_efficient_attention) { - size_t fmha_buffer_bytes = 0; - if (MemoryEfficientAttentionParams::need_workspace(v_head_size, element_size == sizeof(float))) { - fmha_buffer_bytes = batch_size * sequence_length * num_heads * v_head_size * sizeof(float); - } - return qkv_bytes + fmha_buffer_bytes; - } -#else - ORT_UNUSED_PARAMETER(use_memory_efficient_attention); -#endif - - return qkv_bytes + 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length); -} - // Grid: (S, B) // Block: 256 // For unfused PackedAttention @@ -312,7 +258,9 @@ void AddBiasTransposePacked( const int num_heads, const int qk_head_size, const int v_head_size, AttentionQkvFormat format, const int32_t* token_offset, int32_t token_count, cudaStream_t stream) { - if (0 == (qk_head_size & 3) && 0 == (v_head_size & 3)) { + const auto index_width = + GetPackedAttentionQkvMaterializationIndexWidth(qk_head_size, v_head_size); + if (index_width == PackedAttentionQkvMaterializationIndexWidth::Vector4) { using T4Type = typename T4::Type; const int H = qk_head_size / 4; const int H_v = v_head_size / 4; @@ -324,7 +272,7 @@ void AddBiasTransposePacked( batch_size, sequence_length, num_heads, H, H_v, format, token_offset, token_count, stream); - } else if (0 == (qk_head_size & 1) && 0 == (v_head_size & 1)) { + } else if (index_width == PackedAttentionQkvMaterializationIndexWidth::Vector2) { using T2Type = typename T2::Type; const int H = qk_head_size / 2; const int H_v = v_head_size / 2; @@ -460,6 +408,10 @@ Status FusedScaledDotProductAttention( const int v_head_size = parameters.v_head_size; void* fused_runner = data.fused_runner; ORT_RETURN_IF_NOT(nullptr != fused_runner, "fused_runner cannot be NULL"); + ORT_RETURN_IF_NOT( + data.workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::InterleavedTn3h && + data.workspace_recipe.interleaved_qkv_offset_bytes == 0, + "PackedAttention TRT requires a root [T, N, 3, H] QKV workspace."); AddBiasTransposePacked(data.gemm_buffer, data.bias, data.workspace, batch_size, sequence_length, @@ -486,6 +438,8 @@ Status FusedScaledDotProductAttentionCutlass( const int num_heads = parameters.num_heads; const int qk_head_size = parameters.head_size; const int v_head_size = parameters.v_head_size; + ORT_RETURN_IF_NOT(data.workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::Planar, + "PackedAttention MEA requires planar Q/K/V workspace views."); AddBiasTransposePacked(data.gemm_buffer, data.bias, data.workspace, batch_size, sequence_length, num_heads, qk_head_size, v_head_size, @@ -497,15 +451,12 @@ Status FusedScaledDotProductAttentionCutlass( DUMP_TENSOR_D("PackedAttention cutlass data.bias", data.bias, 1, 3 * num_heads * qk_head_size); // Q, K and V pointers - const int model_dimension_qk = num_heads * qk_head_size; - const int model_dimension_v = num_heads * v_head_size; - const size_t elements_qk = static_cast(parameters.token_count) * static_cast(model_dimension_qk); - const size_t elements_v = static_cast(parameters.token_count) * static_cast(model_dimension_v); + const auto& workspace_recipe = data.workspace_recipe; T* qkv = data.workspace; - T* query = qkv; - T* key = query + elements_qk; - T* value = key + elements_qk; - T* accum_workspace = value + elements_v; + T* query = PackedAttentionWorkspaceAt(qkv, workspace_recipe.q_offset_bytes); + T* key = PackedAttentionWorkspaceAt(qkv, workspace_recipe.k_offset_bytes); + T* value = PackedAttentionWorkspaceAt(qkv, workspace_recipe.v_offset_bytes); + T* accum_workspace = PackedAttentionWorkspaceAt(qkv, workspace_recipe.backend_workspace_offset_bytes); DUMP_TENSOR_D("PackedAttention cutlass q(BSNH)", query, parameters.token_count, num_heads * qk_head_size); DUMP_TENSOR_D("PackedAttention cutlass k(BSNH)", key, parameters.token_count, num_heads * qk_head_size); @@ -557,7 +508,6 @@ Status UnfusedScaledDotProductAttention( cudaStream_t stream, PackedAttentionParameters& parameters, PackedAttentionData& data) { - constexpr size_t element_size = sizeof(T); const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; const int num_heads = parameters.num_heads; @@ -565,18 +515,15 @@ Status UnfusedScaledDotProductAttention( const int v_head_size = parameters.v_head_size; const int batches = batch_size * num_heads; - const int size_per_batch_q = sequence_length * qk_head_size; - const int size_per_batch_k = sequence_length * qk_head_size; - const int size_per_batch_v = sequence_length * v_head_size; - const size_t elements_q = static_cast(batches) * static_cast(size_per_batch_q); - const size_t elements_k = static_cast(batches) * static_cast(size_per_batch_k); - const size_t elements_v = static_cast(batches) * static_cast(size_per_batch_v); + const auto& workspace_recipe = data.workspace_recipe; + ORT_RETURN_IF_NOT(workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::Planar, + "PackedAttention unfused attention requires planar Q/K/V workspace views."); // Q, K and V pointers when fused attention is not used T* qkv = data.workspace; - T* q = qkv; - T* k = q + elements_q; - T* v = k + elements_k; + T* q = PackedAttentionWorkspaceAt(qkv, workspace_recipe.q_offset_bytes); + T* k = PackedAttentionWorkspaceAt(qkv, workspace_recipe.k_offset_bytes); + T* v = PackedAttentionWorkspaceAt(qkv, workspace_recipe.v_offset_bytes); AddBiasTransposePacked(data.gemm_buffer, data.bias, data.workspace, batch_size, sequence_length, @@ -584,7 +531,7 @@ Status UnfusedScaledDotProductAttention( AttentionQkvFormat::Q_K_V_BNSH, data.token_offset, parameters.token_count, stream); - T* scaled_qk = qkv + elements_q + elements_k + elements_v; + T* scaled_qk = PackedAttentionWorkspaceAt(qkv, workspace_recipe.backend_workspace_offset_bytes); // Q, K and V are ready now DUMP_TENSOR_INIT(); @@ -613,9 +560,7 @@ Status UnfusedScaledDotProductAttention( DUMP_TENSOR_D("PackedAttention unfused QK", scaled_qk, batch_size * num_heads, sequence_length, sequence_length); - const size_t bytes = GetAttentionScratchSize(element_size, batch_size, num_heads, - sequence_length); - T* attention_score = scaled_qk + (bytes / element_size); + T* attention_score = PackedAttentionWorkspaceAt(qkv, workspace_recipe.second_scratch_offset_bytes); const bool broadcast_attn_bias_dim_0 = parameters.broadcast_attn_bias_dim_0; const bool broadcast_attn_bias_dim_1 = parameters.broadcast_attn_bias_dim_1; diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h index 711718016486f..c1224bf9ef57a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.h @@ -7,29 +7,20 @@ #include #include "contrib_ops/cpu/bert/attention_common.h" #include "contrib_ops/cpu/bert/attention_parameters.h" -#include "contrib_ops/cuda/bert/attention_data.h" +#include "contrib_ops/cuda/bert/packed_attention_data.h" namespace onnxruntime { namespace contrib { namespace cuda { -size_t GetAttentionScratchSize( - size_t element_size, - size_t batch_size, - size_t num_heads, - size_t sequence_length); +template +T* PackedAttentionWorkspaceAt(T* workspace, size_t offset_bytes) { + if (offset_bytes == 0) { + return workspace; + } -size_t GetAttentionWorkspaceSize( - size_t element_size, - size_t batch_size, - size_t num_heads, - size_t qk_head_size, - size_t v_head_size, - size_t sequence_length, - void* fused_runner, - bool use_flash_attention, - bool use_memory_efficient_attention, - bool no_qkv_workspace); + return reinterpret_cast(reinterpret_cast(workspace) + offset_bytes); +} template Status QkvToContext( diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.cc b/onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.cc new file mode 100644 index 0000000000000..e45c30947ca75 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.cc @@ -0,0 +1,1309 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/packed_attention_workspace.h" + +#include + +#include "core/common/safeint.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +namespace { + +constexpr PackedAttentionWorkspaceStatus Ok() noexcept { + return {}; +} + +constexpr PackedAttentionWorkspaceStatus Invalid(const char* message) noexcept { + return {PackedAttentionWorkspaceError::InvalidArgument, message}; +} + +constexpr PackedAttentionWorkspaceStatus Overflow(const char* message) noexcept { + return {PackedAttentionWorkspaceError::Overflow, message}; +} + +PackedAttentionWorkspaceStatus ValidateDimension(int64_t value) noexcept { + if (value < 0) { + return Invalid("Packed attention dimensions must be non-negative."); + } + + if (value > std::numeric_limits::max()) { + return Invalid("Packed attention dimensions and attributes must fit the int32 CUDA ABI."); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus CheckedProductFitsInt32(size_t left, size_t right) noexcept { + size_t product = 0; + auto status = CheckedPackedAttentionMultiply(left, right, product); + if (!status.IsOK()) { + return status; + } + + if (product > static_cast(std::numeric_limits::max())) { + return Invalid("A packed attention derived product does not fit the int32 CUDA ABI."); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus CheckedProductFitsInt32(size_t first, size_t second, size_t third) noexcept { + size_t product = 0; + auto status = CheckedPackedAttentionMultiply(first, second, product); + if (!status.IsOK()) { + return status; + } + + return CheckedProductFitsInt32(product, third); +} + +PackedAttentionWorkspaceStatus CheckedProductFitsInt32(size_t first, size_t second, size_t third, + size_t fourth) noexcept { + size_t product = 0; + auto status = CheckedPackedAttentionMultiply(first, second, product); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionMultiply(product, third, product); + if (!status.IsOK()) { + return status; + } + + return CheckedProductFitsInt32(product, fourth); +} + +PackedAttentionWorkspaceStatus ValidateElementSize(size_t element_size) noexcept { + if (element_size != 2 && element_size != 4) { + return Invalid("Packed attention element size must be 2 or 4 bytes."); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateShape(const PackedAttentionShape& shape, size_t expected_rank) noexcept { + if (shape.rank != expected_rank || shape.rank > shape.dimensions.size()) { + return Invalid("A packed attention input has an invalid rank."); + } + + for (size_t i = 0; i < shape.rank; ++i) { + auto status = ValidateDimension(shape.dimensions[i]); + if (!status.IsOK()) { + return status; + } + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateAttentionBias(const PackedAttentionShape& shape, + int64_t batch_size, + int64_t num_heads, + int64_t sequence_length) noexcept { + if (shape.rank != 4 || shape.rank > shape.dimensions.size()) { + return Invalid("Attention bias must have rank 4."); + } + + for (size_t i = 0; i < shape.rank; ++i) { + auto status = ValidateDimension(shape.dimensions[i]); + if (!status.IsOK()) { + return status; + } + } + + if ((shape.dimensions[0] != 1 && shape.dimensions[0] != batch_size) || + (shape.dimensions[1] != 1 && shape.dimensions[1] != num_heads) || + shape.dimensions[2] != sequence_length || + shape.dimensions[3] != sequence_length) { + return Invalid("Attention bias must have shape [B or 1, N or 1, S, S]."); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateCoreGeometry(int32_t token_count, + int32_t batch_size, + int32_t sequence_length, + int32_t num_heads, + int32_t hidden_size, + int32_t v_hidden_size, + int32_t qk_head_size, + int32_t v_head_size) noexcept { + if (token_count < 0 || batch_size < 0 || sequence_length < 0 || + num_heads <= 0 || hidden_size < 0 || v_hidden_size < 0 || + qk_head_size < 0 || v_head_size < 0) { + return Invalid("Packed attention problem dimensions are invalid."); + } + + const size_t t = static_cast(token_count); + const size_t b = static_cast(batch_size); + const size_t s = static_cast(sequence_length); + const size_t n = static_cast(num_heads); + const size_t h = static_cast(qk_head_size); + const size_t hv = static_cast(v_head_size); + + size_t q_hidden = 0; + auto status = CheckedPackedAttentionMultiply(n, h, q_hidden); + if (!status.IsOK()) { + return status; + } + + size_t v_hidden = 0; + status = CheckedPackedAttentionMultiply(n, hv, v_hidden); + if (!status.IsOK()) { + return status; + } + + if (q_hidden != static_cast(hidden_size) || + v_hidden != static_cast(v_hidden_size)) { + return Invalid("Packed attention hidden sizes do not match the head geometry."); + } + + size_t padded_tokens = 0; + status = CheckedPackedAttentionMultiply(b, s, padded_tokens); + if (!status.IsOK()) { + return status; + } + + if (t > padded_tokens) { + return Invalid("Packed attention token count T must not exceed B * S."); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateQkvMaterializationGeometry(int32_t token_count, + int32_t num_heads, + int32_t qk_head_size, + int32_t v_head_size, + PackedAttentionQkvMaterializationIndexWidth + index_width) noexcept { + const size_t t = static_cast(token_count); + const size_t n = static_cast(num_heads); + const int32_t width = static_cast(index_width); + if (width != 1 && width != 2 && width != 4) { + return Invalid("Packed attention QKV materialization index width is invalid."); + } + + if (qk_head_size % width != 0 || v_head_size % width != 0) { + return Invalid("Packed attention head geometry is not divisible by the QKV materialization index width."); + } + + const size_t h = static_cast(qk_head_size / width); + const size_t hv = static_cast(v_head_size / width); + + size_t qkv_head_size = 0; + auto status = CheckedPackedAttentionAdd(h, h, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAdd(qkv_head_size, hv, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + // The packed transpose producers use int32 offsets in scalar, T2, or T4 + // elements according to the producer selected by the graph adapter. + status = CheckedProductFitsInt32(n, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + return CheckedProductFitsInt32(t, n, qkv_head_size); +} + +PackedAttentionWorkspaceStatus ValidateUnfusedGeometry(int32_t token_count, + int32_t batch_size, + int32_t sequence_length, + int32_t num_heads, + int32_t qk_head_size, + int32_t v_head_size, + PackedAttentionQkvMaterializationIndexWidth + index_width) noexcept { + constexpr int32_t kMaxGridDimY = 65535; + if (batch_size > kMaxGridDimY) { + return Invalid("Packed unfused attention batch size exceeds CUDA gridDim.y."); + } + + const size_t t = static_cast(token_count); + const size_t b = static_cast(batch_size); + const size_t s = static_cast(sequence_length); + const size_t n = static_cast(num_heads); + const size_t h = static_cast(qk_head_size); + const size_t hv = static_cast(v_head_size); + const int32_t width = static_cast(index_width); + if ((width != 1 && width != 2 && width != 4) || + qk_head_size % width != 0 || v_head_size % width != 0) { + return Invalid("Packed attention QKV materialization index width is invalid."); + } + + const size_t producer_h = static_cast(qk_head_size / width); + const size_t producer_hv = static_cast(v_head_size / width); + + auto status = CheckedProductFitsInt32(b, s); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(s, s); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(b, n); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(b, n, s); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(s, h); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(s, hv); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(b, n, s, producer_h); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(b, n, s, producer_hv); + if (!status.IsOK()) { + return status; + } + + size_t qkv_head_size = 0; + status = CheckedPackedAttentionAdd(producer_h, producer_h, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAdd(qkv_head_size, producer_hv, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt32(t, n, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + return CheckedProductFitsInt32(b, s, n, qkv_head_size); +} + +PackedAttentionWorkspaceStatus CheckedProductFitsInt64(size_t first, size_t second, + size_t third) noexcept { + size_t product = 0; + auto status = CheckedPackedAttentionMultiply(first, second, product); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionMultiply(product, third, product); + if (!status.IsOK()) { + return status; + } + + if (product > static_cast(std::numeric_limits::max())) { + return Invalid("A packed attention stride or extent does not fit int64."); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateFusedGrid(int32_t batch_size, int32_t num_heads, + const char* backend_name) noexcept { + constexpr int32_t kMaxGridDimYZ = 65535; + if (batch_size > kMaxGridDimYZ || num_heads > kMaxGridDimYZ) { + return Invalid(backend_name); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateRoundedInt32Dimension(int32_t value, + int32_t alignment, + const char* message) noexcept { + if (value > std::numeric_limits::max() - (alignment - 1)) { + return Invalid(message); + } + + return Ok(); +} + +PackedAttentionWorkspaceStatus ValidateMemoryEfficientGeometry( + size_t element_size, + int32_t batch_size, + int32_t sequence_length, + int32_t num_heads, + int32_t qk_head_size, + int32_t v_head_size, + bool has_attention_bias, + bool broadcast_attn_bias_dim_0, + bool broadcast_attn_bias_dim_1) noexcept { + auto status = ValidateFusedGrid( + batch_size, num_heads, + "Packed memory-efficient attention exceeds CUDA gridDim.y or gridDim.z."); + if (!status.IsOK()) { + return status; + } + + // DispatchBlockSize in fmha_launch_template.h selects at most 64 + // AttentionKernel::kQueriesPerBlock, and AttentionKernel::kAlignLSE is 32. + // Validate both ceil_div round-ups before either CUDA int32 expression is formed. + constexpr int32_t kMaxQueriesPerBlock = 64; + constexpr int32_t kAlignLse = 32; + status = ValidateRoundedInt32Dimension( + sequence_length, kMaxQueriesPerBlock, + "Packed memory-efficient attention query-block rounding exceeds int32."); + if (!status.IsOK()) { + return status; + } + + status = ValidateRoundedInt32Dimension( + sequence_length, kAlignLse, + "Packed memory-efficient attention LSE rounding exceeds int32."); + if (!status.IsOK()) { + return status; + } + + const size_t b = static_cast(batch_size); + const size_t s = static_cast(sequence_length); + const size_t n = static_cast(num_heads); + const size_t h = static_cast(qk_head_size); + const size_t hv = static_cast(v_head_size); + + // CUTLASS stores BSNH batch strides in int64_t. + status = CheckedProductFitsInt64(n, h, s); + if (!status.IsOK()) { + return status; + } + + status = CheckedProductFitsInt64(n, hv, s); + if (!status.IsOK()) { + return status; + } + + if (!has_attention_bias) { + return Ok(); + } + + size_t bias_matrix_elements = 0; + status = CheckedPackedAttentionMultiply(s, s, bias_matrix_elements); + if (!status.IsOK() || + bias_matrix_elements > static_cast(std::numeric_limits::max())) { + return status.IsOK() ? Invalid("The MEA attention-bias head stride does not fit int64.") : status; + } + + const size_t bias_heads = broadcast_attn_bias_dim_1 ? 1 : n; + size_t bias_batch_stride = 0; + status = CheckedPackedAttentionMultiply(bias_heads, bias_matrix_elements, bias_batch_stride); + if (!status.IsOK() || + bias_batch_stride > static_cast(std::numeric_limits::max())) { + return status.IsOK() ? Invalid("The MEA attention-bias batch stride does not fit int64.") : status; + } + + const size_t bias_batches = broadcast_attn_bias_dim_0 ? 1 : b; + size_t bias_extent = 0; + status = CheckedPackedAttentionMultiply(bias_batches, bias_batch_stride, bias_extent); + if (!status.IsOK() || + bias_extent > static_cast(std::numeric_limits::max())) { + return status.IsOK() ? Invalid("The MEA attention-bias extent does not fit int64.") : status; + } + + size_t bias_bytes = 0; + return CheckedPackedAttentionMultiply(bias_extent, element_size, bias_bytes); +} + +PackedAttentionWorkspaceStatus CheckedMultiplyMany(size_t first, size_t second, size_t third, + size_t fourth, size_t fifth, + size_t& result) noexcept { + auto status = CheckedPackedAttentionMultiply(first, second, result); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionMultiply(result, third, result); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionMultiply(result, fourth, result); + if (!status.IsOK()) { + return status; + } + + return CheckedPackedAttentionMultiply(result, fifth, result); +} + +PackedAttentionWorkspaceStatus ComputeWorkspaceRecipe(size_t element_size, + int32_t token_count, + int32_t batch_size, + int32_t sequence_length, + int32_t num_heads, + int32_t qk_head_size, + int32_t v_head_size, + bool has_attention_bias, + bool broadcast_attn_bias_dim_0, + bool broadcast_attn_bias_dim_1, + PackedAttentionBackend backend, + bool no_qkv_workspace, + PackedAttentionQkvMaterializationIndexWidth + qkv_materialization_index_width, + PackedAttentionWorkspaceRecipe& recipe) noexcept { + const size_t t = static_cast(token_count); + const size_t b = static_cast(batch_size); + const size_t s = static_cast(sequence_length); + const size_t n = static_cast(num_heads); + const size_t h = static_cast(qk_head_size); + const size_t hv = static_cast(v_head_size); + + size_t qkv_head_size = 0; + auto status = Ok(); + if (!no_qkv_workspace) { + status = ValidateQkvMaterializationGeometry( + token_count, num_heads, qk_head_size, v_head_size, qkv_materialization_index_width); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAdd(h, h, qkv_head_size); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAdd(qkv_head_size, hv, qkv_head_size); + if (!status.IsOK()) { + return status; + } + } + + switch (backend) { + case PackedAttentionBackend::Trt: + break; + case PackedAttentionBackend::Flash: + status = ValidateFusedGrid( + batch_size, num_heads, + "Packed Flash Attention exceeds CUDA gridDim.y or gridDim.z."); + if (status.IsOK()) { + status = ValidateRoundedInt32Dimension( + sequence_length, 128, + "Packed Flash Attention sequence rounding exceeds int32."); + } + break; + case PackedAttentionBackend::MemoryEfficient: + status = ValidateMemoryEfficientGeometry( + element_size, batch_size, sequence_length, num_heads, qk_head_size, v_head_size, + has_attention_bias, broadcast_attn_bias_dim_0, broadcast_attn_bias_dim_1); + break; + case PackedAttentionBackend::Unfused: + status = ValidateUnfusedGeometry( + token_count, batch_size, sequence_length, num_heads, qk_head_size, v_head_size, + qkv_materialization_index_width); + break; + default: + return Invalid("Packed attention backend is invalid."); + } + + if (!status.IsOK()) { + return status; + } + + size_t qkv_capacity_bytes = 0; + if (!no_qkv_workspace) { + status = CheckedMultiplyMany(element_size, b, s, n, qkv_head_size, qkv_capacity_bytes); + if (!status.IsOK()) { + return status; + } + } + + PackedAttentionWorkspaceRecipe result; + result.no_qkv_workspace = no_qkv_workspace; + result.qkv_capacity_bytes = qkv_capacity_bytes; + + if (!no_qkv_workspace) { + if (backend == PackedAttentionBackend::Trt) { + result.qkv_layout = PackedAttentionQkvWorkspaceLayout::InterleavedTn3h; + result.interleaved_qkv_offset_bytes = 0; + status = CheckedMultiplyMany( + element_size, t, n, qkv_head_size, 1, result.interleaved_qkv_bytes); + if (!status.IsOK()) { + return status; + } + } else { + result.qkv_layout = PackedAttentionQkvWorkspaceLayout::Planar; + result.q_offset_bytes = 0; + + size_t view_tokens = t; + if (backend == PackedAttentionBackend::Unfused) { + status = CheckedPackedAttentionMultiply(b, s, view_tokens); + if (!status.IsOK()) { + return status; + } + } + + status = CheckedMultiplyMany(element_size, view_tokens, n, h, 1, result.q_bytes); + if (!status.IsOK()) { + return status; + } + + // K has the same shape and size as Q. + result.k_offset_bytes = result.q_bytes; + result.k_bytes = result.q_bytes; + status = CheckedPackedAttentionAdd(result.k_offset_bytes, result.k_bytes, result.v_offset_bytes); + if (!status.IsOK()) { + return status; + } + + status = CheckedMultiplyMany(element_size, view_tokens, n, hv, 1, result.v_bytes); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAdd(result.v_offset_bytes, result.v_bytes, + result.backend_workspace_offset_bytes); + if (!status.IsOK()) { + return status; + } + } + } + + switch (backend) { + case PackedAttentionBackend::Trt: + break; + case PackedAttentionBackend::Flash: + // Keep this checked copy in parity with flash::get_softmax_lse_size(S, B, N). + status = CheckedMultiplyMany(sizeof(float), b, s, n, 1, result.backend_workspace_bytes); + if (!status.IsOK()) { + return status; + } + break; + case PackedAttentionBackend::MemoryEfficient: + // Keep this in parity with MemoryEfficientAttentionParams::need_workspace. + if (v_head_size > 128 && element_size != sizeof(float)) { + status = CheckedMultiplyMany(sizeof(float), b, s, n, hv, result.backend_workspace_bytes); + if (!status.IsOK()) { + return status; + } + } + break; + case PackedAttentionBackend::Unfused: { + // Checked equivalent of the dense GetAttentionScratchSize formula. + size_t scratch_bytes = 0; + status = CheckedMultiplyMany(element_size, b, n, s, s, scratch_bytes); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAlign(scratch_bytes, kPackedAttentionWorkspaceAlignment, + result.backend_workspace_bytes); + if (!status.IsOK()) { + return status; + } + + if (result.backend_workspace_offset_bytes != qkv_capacity_bytes) { + return Invalid("Packed unfused attention QKV views do not end at the legacy QKV capacity."); + } + + result.has_second_scratch = true; + status = CheckedPackedAttentionAdd(result.backend_workspace_offset_bytes, + result.backend_workspace_bytes, + result.second_scratch_offset_bytes); + if (!status.IsOK()) { + return status; + } + + size_t both_scratch_bytes = 0; + status = CheckedPackedAttentionMultiply(result.backend_workspace_bytes, 2, both_scratch_bytes); + if (!status.IsOK()) { + return status; + } + + status = CheckedPackedAttentionAdd(qkv_capacity_bytes, both_scratch_bytes, + result.attention_workspace_bytes); + if (!status.IsOK()) { + return status; + } + break; + } + default: + return Invalid("Packed attention backend is invalid."); + } + + if (backend != PackedAttentionBackend::Unfused) { + status = CheckedPackedAttentionAdd(qkv_capacity_bytes, result.backend_workspace_bytes, + result.attention_workspace_bytes); + if (!status.IsOK()) { + return status; + } + } + + status = ValidatePackedAttentionWorkspaceRecipe(result); + if (!status.IsOK()) { + return status; + } + + recipe = result; + return Ok(); +} + +} // namespace + +PackedAttentionWorkspaceStatus CheckedPackedAttentionAdd(size_t left, size_t right, size_t& result) noexcept { + // Keep these helpers non-throwing for the provider/plugin boundary. SafeInt + // construction is intentionally avoided because it reports overflow by throwing. + size_t checked_result = 0; + if (!SafeAdd(left, right, checked_result)) { + return Overflow("Packed attention size addition overflowed size_t."); + } + + result = checked_result; + return Ok(); +} + +PackedAttentionWorkspaceStatus CheckedPackedAttentionMultiply(size_t left, size_t right, + size_t& result) noexcept { + size_t checked_result = 0; + if (!SafeMultiply(left, right, checked_result)) { + return Overflow("Packed attention size multiplication overflowed size_t."); + } + + result = checked_result; + return Ok(); +} + +PackedAttentionWorkspaceStatus CheckedPackedAttentionAlign(size_t value, size_t alignment, + size_t& result) noexcept { + if (alignment == 0) { + return Invalid("Packed attention alignment must be non-zero."); + } + + size_t numerator = 0; + auto status = CheckedPackedAttentionAdd(value, alignment - 1, numerator); + if (!status.IsOK()) { + return status; + } + + const size_t quotient = numerator / alignment; + return CheckedPackedAttentionMultiply(quotient, alignment, result); +} + +PackedAttentionWorkspaceStatus ValidatePackedAttentionWorkspaceRecipe( + const PackedAttentionWorkspaceRecipe& recipe) noexcept { + const auto validate_range = [&recipe](size_t offset, size_t bytes) { + size_t end = 0; + auto status = CheckedPackedAttentionAdd(offset, bytes, end); + if (!status.IsOK()) { + return status; + } + + return end <= recipe.attention_workspace_bytes + ? Ok() + : Invalid("A packed attention workspace view exceeds its allocation."); + }; + + if (recipe.qkv_capacity_bytes > recipe.attention_workspace_bytes) { + return Invalid("Packed attention QKV capacity exceeds its allocation."); + } + + PackedAttentionWorkspaceStatus status; + switch (recipe.qkv_layout) { + case PackedAttentionQkvWorkspaceLayout::None: + if (!recipe.no_qkv_workspace || + recipe.q_offset_bytes != 0 || recipe.q_bytes != 0 || + recipe.k_offset_bytes != 0 || recipe.k_bytes != 0 || + recipe.v_offset_bytes != 0 || recipe.v_bytes != 0 || + recipe.interleaved_qkv_offset_bytes != 0 || recipe.interleaved_qkv_bytes != 0) { + return Invalid("A no-QKV-workspace recipe exposes QKV workspace views."); + } + break; + case PackedAttentionQkvWorkspaceLayout::Planar: + if (recipe.no_qkv_workspace || + recipe.interleaved_qkv_offset_bytes != 0 || recipe.interleaved_qkv_bytes != 0) { + return Invalid("A planar QKV recipe has inconsistent layout fields."); + } + status = validate_range(recipe.q_offset_bytes, recipe.q_bytes); + if (!status.IsOK()) { + return status; + } + status = validate_range(recipe.k_offset_bytes, recipe.k_bytes); + if (!status.IsOK()) { + return status; + } + status = validate_range(recipe.v_offset_bytes, recipe.v_bytes); + if (!status.IsOK()) { + return status; + } + break; + case PackedAttentionQkvWorkspaceLayout::InterleavedTn3h: + if (recipe.no_qkv_workspace || + recipe.q_offset_bytes != 0 || recipe.q_bytes != 0 || + recipe.k_offset_bytes != 0 || recipe.k_bytes != 0 || + recipe.v_offset_bytes != 0 || recipe.v_bytes != 0) { + return Invalid("An interleaved QKV recipe exposes planar Q/K/V views."); + } + status = validate_range(recipe.interleaved_qkv_offset_bytes, recipe.interleaved_qkv_bytes); + if (!status.IsOK()) { + return status; + } + break; + default: + return Invalid("Packed attention QKV workspace layout is invalid."); + } + + status = validate_range(recipe.backend_workspace_offset_bytes, recipe.backend_workspace_bytes); + if (!status.IsOK()) { + return status; + } + + if (recipe.has_second_scratch) { + status = validate_range(recipe.second_scratch_offset_bytes, recipe.backend_workspace_bytes); + if (!status.IsOK()) { + return status; + } + } else if (recipe.second_scratch_offset_bytes != 0) { + return Invalid("A packed attention recipe without a second scratch region has a second-scratch offset."); + } + + return Ok(); +} + +PackedAttentionProblemResult BuildPackedAttentionProblem( + const PackedAttentionInputShapes& inputs) noexcept { + PackedAttentionProblemResult result; + + auto status = ValidateElementSize(inputs.element_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = ValidateDimension(inputs.num_heads); + if (!status.IsOK() || inputs.num_heads == 0) { + result.status = Invalid("PackedAttention num_heads must be positive and fit int32."); + return result; + } + + status = ValidateShape(inputs.input, 2); + if (!status.IsOK()) { + result.status = Invalid("PackedAttention input must have rank 2."); + return result; + } + + status = ValidateShape(inputs.weights, 2); + if (!status.IsOK()) { + result.status = Invalid("PackedAttention weights must have rank 2."); + return result; + } + + status = ValidateShape(inputs.bias, 1); + if (!status.IsOK()) { + result.status = Invalid("PackedAttention bias must have rank 1."); + return result; + } + + status = ValidateShape(inputs.token_offset, 2); + if (!status.IsOK()) { + result.status = Invalid("PackedAttention token_offset must have rank 2."); + return result; + } + + status = ValidateShape(inputs.cumulative_sequence_length, 1); + if (!status.IsOK()) { + result.status = Invalid("PackedAttention cumulative_sequence_length must have rank 1."); + return result; + } + + const int64_t token_count = inputs.input.dimensions[0]; + const int64_t input_hidden_size = inputs.input.dimensions[1]; + const int64_t batch_size = inputs.token_offset.dimensions[0]; + const int64_t sequence_length = inputs.token_offset.dimensions[1]; + const int64_t num_heads = inputs.num_heads; + + if (inputs.weights.dimensions[0] != input_hidden_size) { + result.status = Invalid("PackedAttention weights dimension 0 must equal input hidden size."); + return result; + } + + if (inputs.bias.dimensions[0] != inputs.weights.dimensions[1]) { + result.status = Invalid("PackedAttention bias size must equal weights dimension 1."); + return result; + } + + size_t batch_plus_one = 0; + status = CheckedPackedAttentionAdd(static_cast(batch_size), 1, batch_plus_one); + if (!status.IsOK() || + batch_plus_one > static_cast(std::numeric_limits::max()) || + inputs.cumulative_sequence_length.dimensions[0] != static_cast(batch_plus_one)) { + result.status = Invalid("Cumulative sequence length must have shape [B + 1] within the int32 ABI."); + return result; + } + + int64_t q_hidden_size = 0; + int64_t k_hidden_size = 0; + int64_t v_hidden_size = 0; + if (inputs.qkv_hidden_sizes_count == 0) { + q_hidden_size = inputs.bias.dimensions[0] / 3; + k_hidden_size = q_hidden_size; + v_hidden_size = q_hidden_size; + } else { + if (inputs.qkv_hidden_sizes_count != inputs.qkv_hidden_sizes.size()) { + result.status = Invalid("PackedAttention qkv_hidden_sizes must contain exactly three values."); + return result; + } + + q_hidden_size = inputs.qkv_hidden_sizes[0]; + k_hidden_size = inputs.qkv_hidden_sizes[1]; + v_hidden_size = inputs.qkv_hidden_sizes[2]; + } + + status = ValidateDimension(q_hidden_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = ValidateDimension(k_hidden_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = ValidateDimension(v_hidden_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + if (q_hidden_size != k_hidden_size) { + result.status = Invalid("PackedAttention Q and K hidden sizes must match."); + return result; + } + + if (q_hidden_size % num_heads != 0 || v_hidden_size % num_heads != 0) { + result.status = Invalid("PackedAttention hidden sizes must be divisible by num_heads."); + return result; + } + + size_t qkv_hidden_size = 0; + status = CheckedPackedAttentionAdd(static_cast(q_hidden_size), + static_cast(k_hidden_size), qkv_hidden_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = CheckedPackedAttentionAdd(qkv_hidden_size, static_cast(v_hidden_size), + qkv_hidden_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + if (qkv_hidden_size != static_cast(inputs.bias.dimensions[0]) || + qkv_hidden_size > static_cast(std::numeric_limits::max())) { + result.status = Invalid("PackedAttention Q/K/V hidden sizes must sum to the projection output size."); + return result; + } + + if (inputs.has_attention_bias) { + status = ValidateAttentionBias(inputs.attention_bias, batch_size, num_heads, sequence_length); + if (!status.IsOK()) { + result.status = status; + return result; + } + } + + PackedAttentionProblem problem; + problem.element_size = inputs.element_size; + problem.token_count = static_cast(token_count); + problem.batch_size = static_cast(batch_size); + problem.sequence_length = static_cast(sequence_length); + problem.num_heads = static_cast(num_heads); + problem.input_hidden_size = static_cast(input_hidden_size); + problem.hidden_size = static_cast(q_hidden_size); + problem.v_hidden_size = static_cast(v_hidden_size); + problem.qk_head_size = static_cast(q_hidden_size / num_heads); + problem.v_head_size = static_cast(v_hidden_size / num_heads); + problem.has_attention_bias = inputs.has_attention_bias; + problem.broadcast_attn_bias_dim_0 = + inputs.has_attention_bias && inputs.attention_bias.dimensions[0] == 1; + problem.broadcast_attn_bias_dim_1 = + inputs.has_attention_bias && inputs.attention_bias.dimensions[1] == 1; + problem.qkv_materialization_index_width = + GetPackedAttentionQkvMaterializationIndexWidth(problem.qk_head_size, problem.v_head_size); + + status = ValidateCoreGeometry(problem.token_count, problem.batch_size, problem.sequence_length, + problem.num_heads, problem.hidden_size, problem.v_hidden_size, + problem.qk_head_size, problem.v_head_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + result.problem = problem; + result.status = Ok(); + return result; +} + +PackedAttentionProblemResult BuildPackedMultiHeadAttentionProblem( + const PackedMultiHeadAttentionInputShapes& inputs) noexcept { + PackedAttentionProblemResult result; + + auto status = ValidateElementSize(inputs.element_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = ValidateDimension(inputs.num_heads); + if (!status.IsOK() || inputs.num_heads == 0) { + result.status = Invalid("PackedMultiHeadAttention num_heads must be positive and fit int32."); + return result; + } + + if (inputs.query.rank != 2 && inputs.query.rank != 4) { + result.status = Invalid("PackedMultiHeadAttention query must have rank 2 or 4."); + return result; + } + + status = ValidateShape(inputs.query, inputs.query.rank); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = ValidateShape(inputs.token_offset, 2); + if (!status.IsOK()) { + result.status = Invalid("PackedMultiHeadAttention token_offset must have rank 2."); + return result; + } + + status = ValidateShape(inputs.cumulative_sequence_length, 1); + if (!status.IsOK()) { + result.status = Invalid("PackedMultiHeadAttention cumulative_sequence_length must have rank 1."); + return result; + } + + const int64_t token_count = inputs.query.dimensions[0]; + const int64_t batch_size = inputs.token_offset.dimensions[0]; + const int64_t sequence_length = inputs.token_offset.dimensions[1]; + const int64_t num_heads = inputs.num_heads; + int64_t hidden_size = 0; + int64_t v_hidden_size = 0; + PackedMultiHeadAttentionQkvFormat qkv_format; + + if (inputs.query.rank == 4) { + if (inputs.has_key || inputs.has_value) { + result.status = Invalid("Key and value must be absent when packed QKV is used."); + return result; + } + + if (inputs.query.dimensions[1] != num_heads || + inputs.query.dimensions[2] != 3) { + result.status = Invalid("Packed QKV must have shape [T, N, 3, H]."); + return result; + } + + size_t checked_hidden_size = 0; + status = CheckedPackedAttentionMultiply(static_cast(num_heads), + static_cast(inputs.query.dimensions[3]), + checked_hidden_size); + if (!status.IsOK() || + checked_hidden_size > static_cast(std::numeric_limits::max())) { + result.status = Invalid("Packed QKV hidden size does not fit the int32 CUDA ABI."); + return result; + } + + hidden_size = static_cast(checked_hidden_size); + v_hidden_size = hidden_size; + qkv_format = PackedMultiHeadAttentionQkvFormat::Packed; + } else { + if (!inputs.has_key || !inputs.has_value) { + result.status = Invalid("Separate Q, K, and V inputs must all be present."); + return result; + } + + status = ValidateShape(inputs.key, 2); + if (!status.IsOK()) { + result.status = Invalid("PackedMultiHeadAttention key must have rank 2."); + return result; + } + + status = ValidateShape(inputs.value, 2); + if (!status.IsOK()) { + result.status = Invalid("PackedMultiHeadAttention value must have rank 2."); + return result; + } + + if (inputs.key.dimensions != inputs.query.dimensions) { + result.status = Invalid("Separate query and key shapes must match."); + return result; + } + + if (inputs.value.dimensions[0] != token_count) { + result.status = Invalid("Separate query, key, and value token dimensions must match."); + return result; + } + + hidden_size = inputs.query.dimensions[1]; + v_hidden_size = inputs.value.dimensions[1]; + qkv_format = PackedMultiHeadAttentionQkvFormat::Separate; + } + + if (hidden_size % num_heads != 0 || v_hidden_size % num_heads != 0) { + result.status = Invalid("PackedMultiHeadAttention hidden sizes must be divisible by num_heads."); + return result; + } + + size_t qkv_hidden_size = 0; + status = CheckedPackedAttentionAdd(static_cast(hidden_size), + static_cast(hidden_size), qkv_hidden_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = CheckedPackedAttentionAdd(qkv_hidden_size, static_cast(v_hidden_size), + qkv_hidden_size); + if (!status.IsOK() || + qkv_hidden_size > static_cast(std::numeric_limits::max())) { + result.status = Invalid("PackedMultiHeadAttention Q/K/V hidden-size sum does not fit int32."); + return result; + } + + if (inputs.has_bias) { + status = ValidateShape(inputs.bias, 1); + if (!status.IsOK() || + inputs.bias.dimensions[0] != static_cast(qkv_hidden_size)) { + result.status = Invalid("PackedMultiHeadAttention bias must match the Q/K/V hidden-size sum."); + return result; + } + } + + size_t batch_plus_one = 0; + status = CheckedPackedAttentionAdd(static_cast(batch_size), 1, batch_plus_one); + if (!status.IsOK() || + batch_plus_one > static_cast(std::numeric_limits::max()) || + inputs.cumulative_sequence_length.dimensions[0] != static_cast(batch_plus_one)) { + result.status = Invalid("Cumulative sequence length must have shape [B + 1] within the int32 ABI."); + return result; + } + + if (inputs.has_attention_bias) { + status = ValidateAttentionBias(inputs.attention_bias, batch_size, num_heads, sequence_length); + if (!status.IsOK()) { + result.status = status; + return result; + } + } + + PackedMultiHeadAttentionProblem problem; + problem.element_size = inputs.element_size; + problem.token_count = static_cast(token_count); + problem.batch_size = static_cast(batch_size); + problem.sequence_length = static_cast(sequence_length); + problem.num_heads = static_cast(num_heads); + problem.hidden_size = static_cast(hidden_size); + problem.v_hidden_size = static_cast(v_hidden_size); + problem.qk_head_size = static_cast(hidden_size / num_heads); + problem.v_head_size = static_cast(v_hidden_size / num_heads); + problem.qkv_format = qkv_format; + problem.has_bias = inputs.has_bias; + problem.has_attention_bias = inputs.has_attention_bias; + problem.broadcast_attn_bias_dim_0 = + inputs.has_attention_bias && inputs.attention_bias.dimensions[0] == 1; + problem.broadcast_attn_bias_dim_1 = + inputs.has_attention_bias && inputs.attention_bias.dimensions[1] == 1; + problem.qkv_materialization_index_width = + GetPackedAttentionQkvMaterializationIndexWidth(problem.qk_head_size, problem.v_head_size); + + status = ValidateCoreGeometry(problem.token_count, problem.batch_size, problem.sequence_length, + problem.num_heads, problem.hidden_size, problem.v_hidden_size, + problem.qk_head_size, problem.v_head_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + result.problem = problem; + result.status = Ok(); + return result; +} + +PackedAttentionWorkspaceResult GetPackedAttentionWorkspaceRecipe( + const PackedAttentionProblem& problem) noexcept { + PackedAttentionWorkspaceResult result; + + auto status = ValidateElementSize(problem.element_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = ValidateCoreGeometry(problem.token_count, problem.batch_size, problem.sequence_length, + problem.num_heads, problem.hidden_size, problem.v_hidden_size, + problem.qk_head_size, problem.v_head_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + if (problem.input_hidden_size < 0) { + result.status = Invalid("PackedAttention input hidden size must be non-negative."); + return result; + } + + if (problem.backend == PackedAttentionBackend::Flash) { + result.status = Invalid("PackedAttention does not have a Flash Attention route."); + return result; + } + + if (problem.backend == PackedAttentionBackend::Trt) { + if (!problem.trt_runner_available) { + result.status = Invalid("TRT workspace sizing requires an existing validated runner."); + return result; + } + + if (problem.qk_head_size != problem.v_head_size) { + result.status = Invalid("The packed TRT route requires equal Q/K and V head sizes."); + return result; + } + } + + PackedAttentionWorkspaceRecipe recipe; + status = ComputeWorkspaceRecipe(problem.element_size, problem.token_count, problem.batch_size, + problem.sequence_length, problem.num_heads, problem.qk_head_size, + problem.v_head_size, problem.has_attention_bias, + problem.broadcast_attn_bias_dim_0, problem.broadcast_attn_bias_dim_1, + problem.backend, false, problem.qkv_materialization_index_width, recipe); + if (!status.IsOK()) { + result.status = status; + return result; + } + + size_t projection_n = 0; + status = CheckedPackedAttentionAdd(static_cast(problem.hidden_size), + static_cast(problem.hidden_size), projection_n); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = CheckedPackedAttentionAdd(projection_n, static_cast(problem.v_hidden_size), + projection_n); + if (!status.IsOK() || + projection_n > static_cast(std::numeric_limits::max())) { + result.status = Invalid("PackedAttention projection n does not fit the int32 GEMM ABI."); + return result; + } + + size_t projection_elements = 0; + status = CheckedPackedAttentionMultiply(static_cast(problem.token_count), projection_n, + projection_elements); + if (!status.IsOK()) { + result.status = status; + return result; + } + + status = CheckedPackedAttentionMultiply(projection_elements, problem.element_size, + recipe.projection_bytes); + if (!status.IsOK()) { + result.status = status; + return result; + } + + recipe.projection_m = problem.token_count; + recipe.projection_n = static_cast(projection_n); + recipe.projection_k = problem.input_hidden_size; + + result.recipe = recipe; + result.status = Ok(); + return result; +} + +PackedAttentionWorkspaceResult GetPackedMultiHeadAttentionWorkspaceRecipe( + const PackedMultiHeadAttentionProblem& problem) noexcept { + PackedAttentionWorkspaceResult result; + + auto status = ValidateElementSize(problem.element_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + if (problem.qkv_format != PackedMultiHeadAttentionQkvFormat::Packed && + problem.qkv_format != PackedMultiHeadAttentionQkvFormat::Separate) { + result.status = Invalid("PackedMultiHeadAttention QKV format is invalid."); + return result; + } + + status = ValidateCoreGeometry(problem.token_count, problem.batch_size, problem.sequence_length, + problem.num_heads, problem.hidden_size, problem.v_hidden_size, + problem.qk_head_size, problem.v_head_size); + if (!status.IsOK()) { + result.status = status; + return result; + } + + if (problem.backend == PackedAttentionBackend::Trt) { + if (!problem.trt_runner_available) { + result.status = Invalid("TRT workspace sizing requires an existing validated runner."); + return result; + } + + if (problem.qk_head_size != problem.v_head_size) { + result.status = Invalid("The packed TRT route requires equal Q/K and V head sizes."); + return result; + } + } + + if (problem.backend == PackedAttentionBackend::Flash) { + if (problem.qk_head_size != problem.v_head_size || problem.has_attention_bias) { + result.status = Invalid("The packed Flash route requires equal head sizes and no attention bias."); + return result; + } + } + + const bool no_qkv_workspace = + (problem.backend == PackedAttentionBackend::Trt && + problem.qkv_format == PackedMultiHeadAttentionQkvFormat::Packed && + !problem.has_bias) || + ((problem.backend == PackedAttentionBackend::Flash || + problem.backend == PackedAttentionBackend::MemoryEfficient) && + problem.qkv_format == PackedMultiHeadAttentionQkvFormat::Separate && + !problem.has_bias); + + PackedAttentionWorkspaceRecipe recipe; + status = ComputeWorkspaceRecipe(problem.element_size, problem.token_count, problem.batch_size, + problem.sequence_length, problem.num_heads, problem.qk_head_size, + problem.v_head_size, problem.has_attention_bias, + problem.broadcast_attn_bias_dim_0, problem.broadcast_attn_bias_dim_1, + problem.backend, no_qkv_workspace, + problem.qkv_materialization_index_width, recipe); + if (!status.IsOK()) { + result.status = status; + return result; + } + + result.recipe = recipe; + result.status = Ok(); + return result; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.h b/onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.h new file mode 100644 index 0000000000000..4b9c0ec3fde58 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace.h @@ -0,0 +1,226 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +constexpr size_t kPackedAttentionWorkspaceAlignment = 256; + +enum class PackedAttentionBackend { + Trt, + Flash, + MemoryEfficient, + Unfused, +}; + +enum class PackedAttentionWorkspaceError { + None, + InvalidArgument, + Overflow, +}; + +enum class PackedAttentionQkvWorkspaceLayout { + None, + Planar, + InterleavedTn3h, +}; + +// Width, in scalar elements, of one element indexed by the QKV materialization +// producer. This is a plain host-side description and does not depend on CUDA types. +enum class PackedAttentionQkvMaterializationIndexWidth : int32_t { + Scalar = 1, + Vector2 = 2, + Vector4 = 4, +}; + +constexpr PackedAttentionQkvMaterializationIndexWidth GetPackedAttentionQkvMaterializationIndexWidth( + int32_t qk_head_size, int32_t v_head_size) noexcept { + if (qk_head_size % 4 == 0 && v_head_size % 4 == 0) { + return PackedAttentionQkvMaterializationIndexWidth::Vector4; + } + + if (qk_head_size % 2 == 0 && v_head_size % 2 == 0) { + return PackedAttentionQkvMaterializationIndexWidth::Vector2; + } + + return PackedAttentionQkvMaterializationIndexWidth::Scalar; +} + +struct PackedAttentionWorkspaceStatus { + PackedAttentionWorkspaceError error = PackedAttentionWorkspaceError::None; + const char* message = ""; + + constexpr bool IsOK() const noexcept { + return error == PackedAttentionWorkspaceError::None; + } +}; + +// Packed operator inputs have rank at most four. A rank greater than four is retained +// in rank and rejected without reading beyond dimensions. +struct PackedAttentionShape { + std::array dimensions{}; + size_t rank = 0; +}; + +struct PackedAttentionInputShapes { + PackedAttentionShape input; + PackedAttentionShape weights; + PackedAttentionShape bias; + PackedAttentionShape token_offset; + PackedAttentionShape cumulative_sequence_length; + PackedAttentionShape attention_bias; + size_t element_size = 0; + int64_t num_heads = 0; + size_t qkv_hidden_sizes_count = 0; + std::array qkv_hidden_sizes{}; + bool has_attention_bias = false; +}; + +enum class PackedMultiHeadAttentionQkvFormat { + Packed, + Separate, +}; + +struct PackedMultiHeadAttentionInputShapes { + PackedAttentionShape query; + PackedAttentionShape key; + PackedAttentionShape value; + PackedAttentionShape bias; + PackedAttentionShape token_offset; + PackedAttentionShape cumulative_sequence_length; + PackedAttentionShape attention_bias; + size_t element_size = 0; + int64_t num_heads = 0; + bool has_key = false; + bool has_value = false; + bool has_bias = false; + bool has_attention_bias = false; +}; + +struct PackedAttentionProblem { + size_t element_size = 0; + int32_t token_count = 0; + int32_t batch_size = 0; + int32_t sequence_length = 0; + int32_t num_heads = 0; + int32_t input_hidden_size = 0; + int32_t hidden_size = 0; + int32_t v_hidden_size = 0; + int32_t qk_head_size = 0; + int32_t v_head_size = 0; + bool has_attention_bias = false; + bool broadcast_attn_bias_dim_0 = false; + bool broadcast_attn_bias_dim_1 = false; + PackedAttentionBackend backend = PackedAttentionBackend::Unfused; + bool trt_runner_available = false; + PackedAttentionQkvMaterializationIndexWidth qkv_materialization_index_width = + PackedAttentionQkvMaterializationIndexWidth::Scalar; +}; + +struct PackedMultiHeadAttentionProblem { + size_t element_size = 0; + int32_t token_count = 0; + int32_t batch_size = 0; + int32_t sequence_length = 0; + int32_t num_heads = 0; + int32_t hidden_size = 0; + int32_t v_hidden_size = 0; + int32_t qk_head_size = 0; + int32_t v_head_size = 0; + PackedMultiHeadAttentionQkvFormat qkv_format = PackedMultiHeadAttentionQkvFormat::Packed; + bool has_bias = false; + bool has_attention_bias = false; + bool broadcast_attn_bias_dim_0 = false; + bool broadcast_attn_bias_dim_1 = false; + PackedAttentionBackend backend = PackedAttentionBackend::Unfused; + bool trt_runner_available = false; + PackedAttentionQkvMaterializationIndexWidth qkv_materialization_index_width = + PackedAttentionQkvMaterializationIndexWidth::Scalar; +}; + +template +struct PackedAttentionProblemResult { + PackedAttentionWorkspaceStatus status; + T problem; +}; + +struct PackedAttentionWorkspaceRecipe { + // PackedAttention owns the projection allocation. PackedMultiHeadAttention leaves + // these fields zero because its Q/K/V inputs are already projected. + size_t projection_bytes = 0; + size_t attention_workspace_bytes = 0; + + int32_t projection_m = 0; + int32_t projection_n = 0; + int32_t projection_k = 0; + + // qkv_capacity_bytes preserves the legacy B*S allocation size. It can be + // larger than the route's materialized T-token view. + bool no_qkv_workspace = false; + size_t qkv_capacity_bytes = 0; + + // Planar fields are valid only when qkv_layout is Planar. Q starts at byte + // zero and K has the same byte size as Q. + PackedAttentionQkvWorkspaceLayout qkv_layout = PackedAttentionQkvWorkspaceLayout::None; + size_t q_offset_bytes = 0; + size_t q_bytes = 0; + size_t k_offset_bytes = 0; + size_t k_bytes = 0; + size_t v_offset_bytes = 0; + size_t v_bytes = 0; + + // The TRT materialization producer writes [T, N, 3, H] here. Planar Q/K/V + // offsets are unavailable for this layout. These fields are zero otherwise. + size_t interleaved_qkv_offset_bytes = 0; + size_t interleaved_qkv_bytes = 0; + + // Backend fields are conditional: Flash uses an LSE buffer, MEA may use an + // FP32 accumulator, and unfused uses two equally-sized aligned scratch + // regions. has_second_scratch distinguishes a zero-sized unfused scratch + // region from routes that have no second scratch region. + size_t backend_workspace_offset_bytes = 0; + size_t backend_workspace_bytes = 0; + bool has_second_scratch = false; + size_t second_scratch_offset_bytes = 0; +}; + +struct PackedAttentionWorkspaceResult { + PackedAttentionWorkspaceStatus status; + PackedAttentionWorkspaceRecipe recipe; +}; + +PackedAttentionWorkspaceStatus CheckedPackedAttentionAdd(size_t left, size_t right, size_t& result) noexcept; + +PackedAttentionWorkspaceStatus CheckedPackedAttentionMultiply(size_t left, size_t right, + size_t& result) noexcept; + +PackedAttentionWorkspaceStatus CheckedPackedAttentionAlign(size_t value, size_t alignment, + size_t& result) noexcept; + +// These builders validate input shapes and host-visible geometry only. They do +// not inspect or validate token_offset or cumulative_sequence_length values. +PackedAttentionProblemResult BuildPackedAttentionProblem( + const PackedAttentionInputShapes& inputs) noexcept; + +PackedAttentionProblemResult BuildPackedMultiHeadAttentionProblem( + const PackedMultiHeadAttentionInputShapes& inputs) noexcept; + +PackedAttentionWorkspaceResult GetPackedAttentionWorkspaceRecipe( + const PackedAttentionProblem& problem) noexcept; + +PackedAttentionWorkspaceResult GetPackedMultiHeadAttentionWorkspaceRecipe( + const PackedMultiHeadAttentionProblem& problem) noexcept; + +PackedAttentionWorkspaceStatus ValidatePackedAttentionWorkspaceRecipe( + const PackedAttentionWorkspaceRecipe& recipe) noexcept; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc index 8b159998900ff..5b66b1c35a7c1 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc @@ -10,7 +10,6 @@ #include "contrib_ops/cuda/bert/bert_padding.h" #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" #include "contrib_ops/cuda/bert/flash_attention/flash_api.h" -#include "contrib_ops/cpu/bert/multihead_attention_helper.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -39,7 +38,7 @@ PackedMultiHeadAttention::PackedMultiHeadAttention(const OpKernelInfo& info) : TrtFusedAttention(info) { int64_t num_heads = 0; ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0); - num_heads_ = static_cast(num_heads); + num_heads_ = num_heads; scale_ = info.GetAttrOrDefault("scale", 0.0f); @@ -56,119 +55,48 @@ Status PackedMultiHeadAttention::CheckInputs(const TensorShape& query_shape, const TensorShape& token_offset_shape, const TensorShape& cu_seq_len_shape, const Tensor* attention_bias, - PackedAttentionParameters& parameters) const { - // Shapes of inputs and output: - // When Q, K and V are not packed: - // Input 'query': (token_count, hidden_size) - // Input 'key': (token_count, hidden_size) - // Input 'value': (token_count, v_hidden_size) - // When Q, K and V are packed: - // Input 'query': (token_count, num_heads, 3, head_size) - // Input 'key': None - // Input 'value': None - // Input 'token_offset': (batch_size, sequence_length) - // Input 'cumulative_sequence_length': (batch_size + 1) - // Input 'attention_bias': (batch_size or 1, num_heads or 1, sequence_length, sequence_length) or None - // Output 'output': (token_count, v_hidden_size) - - const auto& query_dims = query_shape.GetDims(); - if (query_dims.size() != 2 && query_dims.size() != 4) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'query' is expected to have 2 or 4 dimensions in packing mode, got ", - query_dims.size()); + PackedAttentionParameters& parameters, + PackedMultiHeadAttentionProblem& problem) const { + PackedMultiHeadAttentionInputShapes inputs; + inputs.query = MakePackedAttentionShape(query_shape); + inputs.token_offset = MakePackedAttentionShape(token_offset_shape); + inputs.cumulative_sequence_length = MakePackedAttentionShape(cu_seq_len_shape); + inputs.element_size = sizeof(T); + inputs.num_heads = GetNumHeads(); + inputs.has_key = key != nullptr; + inputs.has_value = value != nullptr; + inputs.has_bias = bias != nullptr; + inputs.has_attention_bias = attention_bias != nullptr; + if (key != nullptr) { + inputs.key = MakePackedAttentionShape(key->Shape()); } - int64_t token_count = query_dims[0]; - int64_t hidden_size = (query_dims.size() == 2) ? query_dims[1] : (query_dims[1] * query_dims[3]); - - const auto& token_offset_dims = token_offset_shape.GetDims(); - if (token_offset_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'token_offset' is expected to have 2 dimensions in packing mode, got ", - token_offset_dims.size()); + if (value != nullptr) { + inputs.value = MakePackedAttentionShape(value->Shape()); } - - int64_t batch_size = token_offset_dims[0]; - int64_t sequence_length = token_offset_dims[1]; - - int64_t v_hidden_size = hidden_size; - if (query_dims.size() == 4) { - if (key != nullptr || value != nullptr) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'key' and 'value' is expected to be empty when 'query' has 4 dimensions in packing mode"); - } - } else { // query_dims.size() == 2 - if (key == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'key' is expected when 'query' has 2 dimensions in packing mode"); - } - - const auto& key_dims = key->Shape().GetDims(); - if (key_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'key' is expected to have 2 dimension, got ", - key_dims.size()); - } - if (key_dims != query_dims) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' and 'key' is expected to have same shape"); - } - - if (value == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'value' is expected when 'query' has 2 dimensions in packing mode"); - } - const auto& value_dims = value->Shape().GetDims(); - if (value_dims.size() != 2) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value' is expected to have 2 dimensions, got ", - value_dims.size()); - } - if (value_dims[0] != token_count) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 2 dimension 0 should have same length as dimension 0 of input 0"); - } - v_hidden_size = value_dims[1]; - } - if (bias != nullptr) { - const auto& bias_dims = bias->Shape().GetDims(); - if (bias_dims.size() != 1) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' is expected to have 1 dimension, got ", - bias_dims.size()); - } - - if (bias_dims[0] != hidden_size + hidden_size + v_hidden_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'bias' size is expected to be ", - hidden_size + hidden_size + v_hidden_size, ", got ", bias_dims[0]); - } + inputs.bias = MakePackedAttentionShape(bias->Shape()); } - - const auto& cu_seq_len_dims = cu_seq_len_shape.GetDims(); - if (cu_seq_len_dims.size() != 1 || cu_seq_len_dims[0] != batch_size + 1) { - return ORT_MAKE_STATUS( - ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'cumulative_sequence_length' should have 1 dimension with size equal to batch_size + 1"); + if (attention_bias != nullptr) { + inputs.attention_bias = MakePackedAttentionShape(attention_bias->Shape()); } - const int num_heads = this->GetNumHeads(); + auto problem_result = BuildPackedMultiHeadAttentionProblem(inputs); + ORT_RETURN_IF_ERROR(PackedAttentionWorkspaceStatusToStatus(problem_result.status)); + problem = problem_result.problem; - gsl::span attention_bias_dims; - if (attention_bias != nullptr) { - attention_bias_dims = attention_bias->Shape().GetDims(); - ORT_RETURN_IF_ERROR(multihead_attention_helper::CheckAttentionBias( - attention_bias_dims, batch_size, num_heads, sequence_length, sequence_length)); - } - parameters.broadcast_attn_bias_dim_0 = attention_bias_dims.size() > 0 && attention_bias_dims[0] == 1; - parameters.broadcast_attn_bias_dim_1 = attention_bias_dims.size() > 1 && attention_bias_dims[1] == 1; + parameters.broadcast_attn_bias_dim_0 = problem.broadcast_attn_bias_dim_0; + parameters.broadcast_attn_bias_dim_1 = problem.broadcast_attn_bias_dim_1; - parameters.batch_size = static_cast(batch_size); - parameters.sequence_length = static_cast(sequence_length); + parameters.batch_size = problem.batch_size; + parameters.sequence_length = problem.sequence_length; parameters.input_hidden_size = -1; // not applicable - parameters.hidden_size = static_cast(hidden_size); - parameters.v_hidden_size = static_cast(v_hidden_size); - parameters.head_size = static_cast(hidden_size) / num_heads; - parameters.v_head_size = static_cast(v_hidden_size) / num_heads; - parameters.num_heads = num_heads; + parameters.hidden_size = problem.hidden_size; + parameters.v_hidden_size = problem.v_hidden_size; + parameters.head_size = problem.qk_head_size; + parameters.v_head_size = problem.v_head_size; + parameters.num_heads = problem.num_heads; parameters.scale = this->GetScale(); - parameters.token_count = static_cast(token_count); + parameters.token_count = problem.token_count; return Status::OK(); } @@ -186,6 +114,7 @@ Status PackedMultiHeadAttention::ComputeInternal(OpKernelContext* context) co typedef typename ToCudaType::MappedType CudaT; PackedAttentionParameters parameters; + PackedMultiHeadAttentionProblem problem; parameters.use_tf32 = this->UseTF32(); ORT_RETURN_IF_ERROR(CheckInputs(query->Shape(), key, @@ -194,7 +123,8 @@ Status PackedMultiHeadAttention::ComputeInternal(OpKernelContext* context) co token_offset->Shape(), cumulative_sequence_length->Shape(), attention_bias, - parameters)); + parameters, + problem)); TensorShapeVector output_shape{parameters.token_count, parameters.v_hidden_size}; Tensor* output = context->Output(0, output_shape); @@ -251,23 +181,20 @@ Status PackedMultiHeadAttention::ComputeInternal(OpKernelContext* context) co cublasHandle_t cublas = this->GetCublasHandle(context); - constexpr size_t element_size = sizeof(T); - // When the source and target format is same (like TN3H => TN3H, or TNH => TNH) and no bias, need not transpose qkv. - const bool no_qkv_workspace = (fused_runner != nullptr && key == nullptr && bias == nullptr) || - ((use_memory_efficient_attention || use_flash_attention) && - value != nullptr && - bias == nullptr); - size_t workSpaceSize = GetAttentionWorkspaceSize(element_size, - parameters.batch_size, - parameters.num_heads, - parameters.head_size, - parameters.v_head_size, - parameters.sequence_length, - fused_runner, - use_flash_attention, - use_memory_efficient_attention, - no_qkv_workspace); - auto work_space = this->template GetScratchBuffer(workSpaceSize, this->GetComputeStream(context)); + problem.backend = use_flash_attention + ? PackedAttentionBackend::Flash + : (fused_runner != nullptr + ? PackedAttentionBackend::Trt + : (use_memory_efficient_attention + ? PackedAttentionBackend::MemoryEfficient + : PackedAttentionBackend::Unfused)); + problem.trt_runner_available = fused_runner != nullptr; + auto workspace_result = GetPackedMultiHeadAttentionWorkspaceRecipe(problem); + ORT_RETURN_IF_ERROR(PackedAttentionWorkspaceStatusToStatus(workspace_result.status)); + const PackedAttentionWorkspaceRecipe& workspace_recipe = workspace_result.recipe; + + auto work_space = this->template GetScratchBuffer( + workspace_recipe.attention_workspace_bytes, this->GetComputeStream(context)); PackedMultiHeadAttentionData data; data.query = reinterpret_cast(query->Data()); @@ -284,8 +211,12 @@ Status PackedMultiHeadAttention::ComputeInternal(OpKernelContext* context) co data.fused_runner = reinterpret_cast(fused_runner); data.use_flash_attention = use_flash_attention; data.use_memory_efficient_attention = use_memory_efficient_attention; - data.no_qkv_workspace = no_qkv_workspace; - data.source_qkv_format = (key == nullptr) ? AttentionQkvFormat::QKV_TN3H : AttentionQkvFormat::Q_K_V_TNH; + data.no_qkv_workspace = workspace_recipe.no_qkv_workspace; + data.source_qkv_format = + problem.qkv_format == PackedMultiHeadAttentionQkvFormat::Packed + ? AttentionQkvFormat::QKV_TN3H + : AttentionQkvFormat::Q_K_V_TNH; + data.workspace_recipe = workspace_recipe; return QkvToContext(device_prop, cublas, this->Stream(context), parameters, data); } diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.h b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.h index 3e59ce3dd229e..6b9ee629fde1d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.h @@ -24,12 +24,13 @@ class PackedMultiHeadAttention final : public TrtFusedAttention { const TensorShape& token_offset_shape, const TensorShape& cu_seq_len_shape, const Tensor* attention_bias, - PackedAttentionParameters& parameters) const; - int GetNumHeads() const { return num_heads_; } + PackedAttentionParameters& parameters, + PackedMultiHeadAttentionProblem& problem) const; + int64_t GetNumHeads() const { return num_heads_; } float GetScale() const { return scale_; } - int num_heads_; // number of attention heads - float scale_; // the scale for softmax in memory efficient attention or unfused attention. + int64_t num_heads_; // number of attention heads + float scale_; // the scale for softmax in memory efficient attention or unfused attention. bool disable_memory_efficient_attention_; bool disable_flash_attention_; diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu index db70ad9b8b064..74f52a6152318 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu @@ -518,7 +518,9 @@ void AddBiasTransposePacked( AttentionQkvFormat source_format, AttentionQkvFormat target_format, const int32_t* token_offset, int32_t token_count, cudaStream_t stream) { - if (0 == (qk_head_size & 3) && 0 == (v_head_size & 3)) { + const auto index_width = + GetPackedAttentionQkvMaterializationIndexWidth(qk_head_size, v_head_size); + if (index_width == PackedAttentionQkvMaterializationIndexWidth::Vector4) { using T4Type = typename T4::Type; const int H = qk_head_size / 4; const int H_v = v_head_size / 4; @@ -533,7 +535,7 @@ void AddBiasTransposePacked( num_heads, H, H_v, source_format, target_format, token_offset, token_count, stream); - } else if (0 == (qk_head_size & 1) && 0 == (v_head_size & 1)) { + } else if (index_width == PackedAttentionQkvMaterializationIndexWidth::Vector2) { using T2Type = typename T2::Type; const int H = qk_head_size / 2; const int H_v = v_head_size / 2; @@ -571,6 +573,13 @@ Status FusedAttentionTrt( const int v_head_size = parameters.v_head_size; void* fused_runner = data.fused_runner; ORT_RETURN_IF_NOT(nullptr != fused_runner, "fused_runner cannot be NULL"); + ORT_RETURN_IF_NOT( + (data.no_qkv_workspace && + data.workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::None) || + (!data.no_qkv_workspace && + data.workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::InterleavedTn3h && + data.workspace_recipe.interleaved_qkv_offset_bytes == 0), + "PackedMultiHeadAttention TRT requires direct or root [T, N, 3, H] QKV."); // When packed QKV is used, we can directly pass it to fused runner. Otherwise, we need transpose to BSN3H format. const T* qkv = data.query; @@ -603,11 +612,13 @@ Status FlashAttention( const int qk_head_size = parameters.head_size; const int v_head_size = parameters.v_head_size; - // Q, K and V pointers - const int model_dimension_qk = num_heads * qk_head_size; - const int model_dimension_v = num_heads * v_head_size; - const size_t elements_qk = static_cast(parameters.token_count) * static_cast(model_dimension_qk); - const size_t elements_v = static_cast(parameters.token_count) * static_cast(model_dimension_v); + const auto& workspace_recipe = data.workspace_recipe; + ORT_RETURN_IF_NOT( + (data.no_qkv_workspace && + workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::None) || + (!data.no_qkv_workspace && + workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::Planar), + "PackedMultiHeadAttention Flash requires direct or planar Q/K/V views."); // When separated Q, K, V is used, we can directly use them in Cutlass FMHA. Otherwise, transpose BSN3H to 3BSNH if (!data.no_qkv_workspace) { @@ -622,12 +633,19 @@ Status FlashAttention( : parameters.scale; int32_t* cu_seqlens_q = const_cast(data.cumulative_sequence_length); int32_t* cu_seqlens_k = const_cast(data.cumulative_sequence_length); - const void* query = data.no_qkv_workspace ? data.query : data.workspace; - const void* key = data.no_qkv_workspace ? data.key : (data.workspace + elements_qk); - const void* value = data.no_qkv_workspace ? data.value : (data.workspace + elements_qk + elements_qk); + const void* query = data.no_qkv_workspace + ? data.query + : PackedAttentionWorkspaceAt(data.workspace, workspace_recipe.q_offset_bytes); + const void* key = data.no_qkv_workspace + ? data.key + : PackedAttentionWorkspaceAt(data.workspace, workspace_recipe.k_offset_bytes); + const void* value = data.no_qkv_workspace + ? data.value + : PackedAttentionWorkspaceAt(data.workspace, workspace_recipe.v_offset_bytes); void* softmax_lse_buffer = data.no_qkv_workspace ? data.workspace - : (data.workspace + elements_qk + elements_qk + elements_v); + : PackedAttentionWorkspaceAt( + data.workspace, workspace_recipe.backend_workspace_offset_bytes); ORT_RETURN_IF_ERROR( onnxruntime::flash::mha_varlen_fwd( @@ -679,11 +697,13 @@ Status FusedAttentionCutlass( const int qk_head_size = parameters.head_size; const int v_head_size = parameters.v_head_size; - // Q, K and V pointers - const int model_dimension_qk = num_heads * qk_head_size; - const int model_dimension_v = num_heads * v_head_size; - const size_t elements_qk = static_cast(parameters.token_count) * static_cast(model_dimension_qk); - const size_t elements_v = static_cast(parameters.token_count) * static_cast(model_dimension_v); + const auto& workspace_recipe = data.workspace_recipe; + ORT_RETURN_IF_NOT( + (data.no_qkv_workspace && + workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::None) || + (!data.no_qkv_workspace && + workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::Planar), + "PackedMultiHeadAttention MEA requires direct or planar Q/K/V views."); // When separated Q, K, V is used, we can directly use them in Cutlass FMHA. Otherwise, transpose BSN3H to 3BSNH if (!data.no_qkv_workspace) { @@ -712,9 +732,15 @@ Status FusedAttentionCutlass( p.seqlen_k_ptr = nullptr; p.seqstart_q_ptr = data.cumulative_sequence_length; p.seqstart_k_ptr = data.cumulative_sequence_length; - p.query = data.no_qkv_workspace ? data.query : data.workspace; - p.key = data.no_qkv_workspace ? data.key : (data.workspace + elements_qk); - p.value = data.no_qkv_workspace ? data.value : (data.workspace + elements_qk + elements_qk); + p.query = data.no_qkv_workspace + ? data.query + : PackedAttentionWorkspaceAt(data.workspace, workspace_recipe.q_offset_bytes); + p.key = data.no_qkv_workspace + ? data.key + : PackedAttentionWorkspaceAt(data.workspace, workspace_recipe.k_offset_bytes); + p.value = data.no_qkv_workspace + ? data.value + : PackedAttentionWorkspaceAt(data.workspace, workspace_recipe.v_offset_bytes); p.attn_bias = data.attention_bias; p.broadcast_attn_bias_dim_0 = parameters.broadcast_attn_bias_dim_0; @@ -723,7 +749,8 @@ Status FusedAttentionCutlass( p.output = data.output; p.is_kv_bsnh = true; p.workspace = MemoryEfficientAttentionParams::need_workspace(v_head_size, sizeof(T) == sizeof(float)) - ? (data.workspace + (data.no_qkv_workspace ? 0 : (elements_qk + elements_qk + elements_v))) + ? PackedAttentionWorkspaceAt( + data.workspace, workspace_recipe.backend_workspace_offset_bytes) : nullptr; p.stream = stream; p.has_custom_right_padding = false; @@ -747,7 +774,6 @@ Status UnfusedAttention( cudaStream_t stream, PackedAttentionParameters& parameters, PackedMultiHeadAttentionData& data) { - constexpr size_t element_size = sizeof(T); const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; const int num_heads = parameters.num_heads; @@ -755,12 +781,9 @@ Status UnfusedAttention( const int v_head_size = parameters.v_head_size; const int batches = batch_size * num_heads; - const int size_per_batch_q = sequence_length * qk_head_size; - const int size_per_batch_k = sequence_length * qk_head_size; - const int size_per_batch_v = sequence_length * v_head_size; - const size_t elements_q = static_cast(batches) * static_cast(size_per_batch_q); - const size_t elements_k = static_cast(batches) * static_cast(size_per_batch_k); - const size_t elements_v = static_cast(batches) * static_cast(size_per_batch_v); + const auto& workspace_recipe = data.workspace_recipe; + ORT_RETURN_IF_NOT(workspace_recipe.qkv_layout == PackedAttentionQkvWorkspaceLayout::Planar, + "PackedMultiHeadAttention unfused attention requires planar Q/K/V workspace views."); // Q, K and V pointers when fused attention is not used AddBiasTransposePacked(data.query, data.key, data.value, data.bias, data.workspace, @@ -770,10 +793,10 @@ Status UnfusedAttention( data.token_offset, parameters.token_count, stream); T* qkv = data.workspace; - T* q = qkv; - T* k = q + elements_q; - T* v = k + elements_k; - T* scaled_qk = qkv + elements_q + elements_k + elements_v; + T* q = PackedAttentionWorkspaceAt(qkv, workspace_recipe.q_offset_bytes); + T* k = PackedAttentionWorkspaceAt(qkv, workspace_recipe.k_offset_bytes); + T* v = PackedAttentionWorkspaceAt(qkv, workspace_recipe.v_offset_bytes); + T* scaled_qk = PackedAttentionWorkspaceAt(qkv, workspace_recipe.backend_workspace_offset_bytes); // Compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scaled_qk: BxNxSxT // Q: BxNxSxH, K: BxNxSxH, Q*K': BxNxSxS @@ -801,9 +824,7 @@ Status UnfusedAttention( DUMP_TENSOR_D("v (BNSH)", v, batch_size, num_heads, sequence_length, v_head_size); DUMP_TENSOR_D("QK", scaled_qk, batch_size, num_heads, sequence_length, sequence_length); - const size_t bytes = GetAttentionScratchSize(element_size, batch_size, num_heads, - sequence_length); - T* attention_score = scaled_qk + (bytes / element_size); + T* attention_score = PackedAttentionWorkspaceAt(qkv, workspace_recipe.second_scratch_offset_bytes); // Apply softmax and store result R to attention_score: BxNxSxS ORT_RETURN_IF_ERROR(ComputeSoftmaxWithCumSeqLength( diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.h index 10e95b95657f1..7019962286809 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.h @@ -7,7 +7,7 @@ #include #include "contrib_ops/cpu/bert/attention_common.h" #include "contrib_ops/cpu/bert/attention_parameters.h" -#include "contrib_ops/cuda/bert/attention_data.h" +#include "contrib_ops/cuda/bert/packed_attention_data.h" namespace onnxruntime { namespace contrib { diff --git a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc index 96c629b4616d5..74122001bd796 100644 --- a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc @@ -15,6 +15,80 @@ namespace onnxruntime { using contrib::AttentionMaskType; namespace test { +namespace { + +class ScopedStdoutCapture { + public: + explicit ScopedStdoutCapture(bool enabled) : capturing_(enabled) { + if (capturing_) { + testing::internal::CaptureStdout(); + } + } + + ~ScopedStdoutCapture() { + if (capturing_) { + (void)testing::internal::GetCapturedStdout(); + } + } + + std::string Stop() { + capturing_ = false; + return testing::internal::GetCapturedStdout(); + } + + private: + bool capturing_; +}; + +template +void RunAndVerifyAttentionRoute(const char* expected_route, Run&& run) { + ScopedStdoutCapture capture(true); + run(); + const std::string debug_output = capture.Stop(); + EXPECT_NE(debug_output.find(expected_route), std::string::npos) << debug_output; +} + +bool IsTrtFusedAttentionRouteObservable(int head_size, int sequence_length) { +#if USE_TRT_FUSED_ATTENTION + if (!HasCudaEnvironment(0)) { + return false; + } + + // These tests use the non-flash branch of FusedMHARunnerFP16v2::IsSupported. + // Keep its SM allowlist, head-size rules, and sequence cap in sync with mha_runner.cu. + const int sm = GetCudaArchitecture() / 10; + const bool supported_sm = sm == 70 || sm == 75 || sm == 80 || + sm == 86 || sm == 89; + return supported_sm && + (head_size == 32 || head_size == 64) && + !(sm == 70 && head_size == 32) && + sequence_length <= 384; +#else + ORT_UNUSED_PARAMETER(head_size); + ORT_UNUSED_PARAMETER(sequence_length); + return false; +#endif +} + +bool IsMemoryEfficientAttentionRouteObservable(int head_size) { +#if USE_MEMORY_EFFICIENT_ATTENTION + return HasCudaEnvironment(530) && + head_size % 8 == 0 && + head_size <= 1024; +#else + ORT_UNUSED_PARAMETER(head_size); + return false; +#endif +} + +enum class PackedAttentionRoute { + Trt, + MemoryEfficient, + Unfused, +}; + +} // namespace + static void RunPackedAttentionTest( const std::vector& input_data, // input: [token_count, hidden_size] const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] @@ -94,6 +168,83 @@ static void RunPackedAttentionTest( } } +static void RunPackedAttentionRouteTest(PackedAttentionRoute route) { + constexpr int kBatchSize = 2; + constexpr int kSequenceLength = 2; + constexpr int kTokenCount = 3; + constexpr int kHiddenSize = 32; + constexpr int kNumHeads = 1; + constexpr int kHeadSize = kHiddenSize / kNumHeads; + static_assert(kTokenCount < kBatchSize * kSequenceLength); + + const std::vector input_data(kTokenCount * kHiddenSize, 0.0f); + const std::vector weight_data(kHiddenSize * 3 * kHiddenSize, 0.0f); + const std::vector bias_data(3 * kHiddenSize, 0.0f); + const std::vector token_offset{0, 2, 3, 1}; + const std::vector cumulative_sequence_length{0, 1, 3}; + const std::vector output_data(kTokenCount * kHiddenSize, 0.0f); + + if (route == PackedAttentionRoute::Trt) { + if (!IsTrtFusedAttentionRouteObservable(kHeadSize, kSequenceLength)) { + GTEST_SKIP() << "PackedAttention TRT route is unavailable in this build or CUDA environment."; + } + + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "0"}, + {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "0"}, + {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + RunAndVerifyAttentionRoute("SdpaKernel=TRT_FUSED_ATTENTION", [&]() { + RunPackedAttentionTest( + input_data, weight_data, bias_data, token_offset, cumulative_sequence_length, + output_data, kBatchSize, kSequenceLength, kHiddenSize, kNumHeads, kTokenCount, + true, true, {}, {}); + }); + return; + } + + if (route == PackedAttentionRoute::MemoryEfficient) { + if (!IsMemoryEfficientAttentionRouteObservable(kHeadSize)) { + GTEST_SKIP() << "PackedAttention MEA route is unavailable in this build or CUDA environment."; + } + + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + RunAndVerifyAttentionRoute("SdpaKernel=EFFICIENT_ATTENTION", [&]() { + RunPackedAttentionTest( + input_data, weight_data, bias_data, token_offset, cumulative_sequence_length, + output_data, kBatchSize, kSequenceLength, kHiddenSize, kNumHeads, kTokenCount, + true, true, {}, {}); + }); + return; + } + + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "PackedAttention MATH route requires a CUDA device."; + } + + // PackedAttention does not honor the MEA-disable option. FP32 is deliberately + // MEA-ineligible here, so this invocation deterministically + // observes the unfused route without changing runtime behavior. + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{ + {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + RunAndVerifyAttentionRoute("SdpaKernel=MATH", [&]() { + RunPackedAttentionTest( + input_data, weight_data, bias_data, token_offset, cumulative_sequence_length, + output_data, kBatchSize, kSequenceLength, kHiddenSize, kNumHeads, kTokenCount, + false, true, {}, {}); + }); +} + static void RunPackedAttentionTest( const std::vector& input_data, // input: [token_count, hidden_size] const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] @@ -132,6 +283,18 @@ static void RunPackedAttentionTest( InvokePackedAttentionTest(false, false); } +TEST(PackedAttentionTest, PackedRouteObservedTrtWithPadding) { + RunPackedAttentionRouteTest(PackedAttentionRoute::Trt); +} + +TEST(PackedAttentionTest, PackedRouteObservedMemoryEfficientWithPadding) { + RunPackedAttentionRouteTest(PackedAttentionRoute::MemoryEfficient); +} + +TEST(PackedAttentionTest, PackedRouteObservedUnfusedWithPadding) { + RunPackedAttentionRouteTest(PackedAttentionRoute::Unfused); +} + TEST(PackedAttentionTest, NoPack) { int batch_size = 1; int sequence_length = 2; diff --git a/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc index 17862c0aca6fa..5168f007c6e57 100644 --- a/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc @@ -15,6 +15,125 @@ namespace onnxruntime { using contrib::AttentionMaskType; namespace test { +namespace { + +class ScopedStdoutCapture { + public: + explicit ScopedStdoutCapture(bool enabled) : capturing_(enabled) { + if (capturing_) { + testing::internal::CaptureStdout(); + } + } + + ~ScopedStdoutCapture() { + if (capturing_) { + (void)testing::internal::GetCapturedStdout(); + } + } + + std::string Stop() { + capturing_ = false; + return testing::internal::GetCapturedStdout(); + } + + private: + bool capturing_; +}; + +template +void RunAndVerifyAttentionRoute(const char* expected_route, Run&& run) { + ScopedStdoutCapture capture(true); + run(); + const std::string debug_output = capture.Stop(); + EXPECT_NE(debug_output.find(expected_route), std::string::npos) << debug_output; +} + +bool IsTrtFusedAttentionRouteObservable(int qk_head_size, + int v_head_size, + int sequence_length, + bool has_attention_bias) { + if (qk_head_size != v_head_size || has_attention_bias) { + return false; + } + +#if USE_TRT_FUSED_ATTENTION + if (!HasCudaEnvironment(0)) { + return false; + } + + // These tests use the non-flash branch of FusedMHARunnerFP16v2::IsSupported. + // Keep its SM allowlist, head-size rules, and sequence cap in sync with mha_runner.cu. + const int sm = GetCudaArchitecture() / 10; + const bool supported_sm = sm == 70 || sm == 75 || sm == 80 || + sm == 86 || sm == 89; + return supported_sm && + (qk_head_size == 32 || qk_head_size == 64) && + !(sm == 70 && qk_head_size == 32) && + sequence_length <= 384; +#else + ORT_UNUSED_PARAMETER(qk_head_size); + ORT_UNUSED_PARAMETER(v_head_size); + ORT_UNUSED_PARAMETER(sequence_length); + ORT_UNUSED_PARAMETER(has_attention_bias); + return false; +#endif +} + +bool IsMemoryEfficientAttentionGeometrySupported(int qk_head_size, + int v_head_size, + int sequence_length, + bool has_attention_bias) { + return qk_head_size % 8 == 0 && + v_head_size % 8 == 0 && + qk_head_size <= 1024 && + v_head_size <= 1024 && + (!has_attention_bias || sequence_length % (4 * sizeof(MLFloat16)) == 0); +} + +bool IsMemoryEfficientAttentionRouteObservable(int qk_head_size, + int v_head_size, + int sequence_length, + bool has_attention_bias) { +#if USE_MEMORY_EFFICIENT_ATTENTION + return HasCudaEnvironment(530) && + IsMemoryEfficientAttentionGeometrySupported( + qk_head_size, v_head_size, sequence_length, has_attention_bias); +#else + ORT_UNUSED_PARAMETER(qk_head_size); + ORT_UNUSED_PARAMETER(v_head_size); + ORT_UNUSED_PARAMETER(sequence_length); + ORT_UNUSED_PARAMETER(has_attention_bias); + return false; +#endif +} + +bool IsFlashAttentionRouteObservable(int qk_head_size, + int v_head_size, + bool has_attention_bias) { +#if USE_FLASH_ATTENTION + if (!HasCudaEnvironment(800) || + qk_head_size != v_head_size || + qk_head_size % 8 != 0 || + qk_head_size > 256 || + has_attention_bias) { + return false; + } + +#ifdef ORT_QUICK_BUILD + return qk_head_size == 128; +#else + return true; +#endif +#else + ORT_UNUSED_PARAMETER(qk_head_size); + ORT_UNUSED_PARAMETER(v_head_size); + ORT_UNUSED_PARAMETER(has_attention_bias); + return false; +#endif +} + +} // namespace + #define InvokePackedMultiHeadAttentionTest(use_float16, use_scale) \ RunPackedMultiHeadAttentionTest( \ query_data, \ @@ -159,53 +278,101 @@ static void RunPackedMultiHeadAttentionTest( AttentionKernelType kernel_type, const std::vector& attention_bias_data = {}, bool broadcast_attention_bias = false) { + const int qk_head_size = hidden_size / number_of_heads; + const int v_head_size = v_hidden_size / number_of_heads; + const bool has_attention_bias = !attention_bias_data.empty(); + if (kernel_type == AttentionKernelType::AttentionKernel_TrtFusedAttention) { + if (!IsTrtFusedAttentionRouteObservable( + qk_head_size, v_head_size, sequence_length, has_attention_bias)) { + GTEST_SKIP() << "PackedMultiHeadAttention TRT route is unavailable for this configuration."; + } + ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ {onnxruntime::contrib::attention::kDisableFlashAttention, "1"}, {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "0"}, {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "0"}, {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, - {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}}; - InvokePackedMultiHeadAttentionTest(true, true); - InvokePackedMultiHeadAttentionTest(true, false); + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + RunAndVerifyAttentionRoute("SdpaKernel=TRT_FUSED_ATTENTION", [&]() { + InvokePackedMultiHeadAttentionTest(true, true); + }); + RunAndVerifyAttentionRoute("SdpaKernel=TRT_FUSED_ATTENTION", [&]() { + InvokePackedMultiHeadAttentionTest(true, false); + }); } -#if USE_MEMORY_EFFICIENT_ATTENTION if (kernel_type == AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention) { + const bool geometry_supports_mea = IsMemoryEfficientAttentionGeometrySupported( + qk_head_size, v_head_size, sequence_length, has_attention_bias); + if (geometry_supports_mea && + !IsMemoryEfficientAttentionRouteObservable( + qk_head_size, v_head_size, sequence_length, has_attention_bias)) { + GTEST_SKIP() << "PackedMultiHeadAttention MEA route is unavailable in this build or CUDA environment."; + } + + if (!geometry_supports_mea && !HasCudaEnvironment(530)) { + GTEST_SKIP() << "PackedMultiHeadAttention invalid-geometry fallback requires a CUDA device."; + } + ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ {onnxruntime::contrib::attention::kDisableFlashAttention, "1"}, {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, - {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "0"}}}; - InvokePackedMultiHeadAttentionTest(true, true); - InvokePackedMultiHeadAttentionTest(true, false); - // Cutlass FMHA need sequence length >= 256 to trigger, so we only test fp16 here. + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "0"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + const char* expected_route = + geometry_supports_mea ? "SdpaKernel=EFFICIENT_ATTENTION" : "SdpaKernel=MATH"; + RunAndVerifyAttentionRoute(expected_route, [&]() { + InvokePackedMultiHeadAttentionTest(true, true); + }); + RunAndVerifyAttentionRoute(expected_route, [&]() { + InvokePackedMultiHeadAttentionTest(true, false); + }); } -#endif -#if USE_FLASH_ATTENTION if (kernel_type == AttentionKernelType::AttentionKernel_FlashAttention) { + if (!IsFlashAttentionRouteObservable(qk_head_size, v_head_size, has_attention_bias)) { + GTEST_SKIP() << "PackedMultiHeadAttention Flash route is unavailable in this build or CUDA environment."; + } + ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ {onnxruntime::contrib::attention::kDisableFlashAttention, "0"}, - {onnxruntime::contrib::attention::kMinSeqLenForFlashAttentionPackedQKV, "0"}}}; - InvokePackedMultiHeadAttentionTest(true, true); + {onnxruntime::contrib::attention::kMinSeqLenForFlashAttentionPackedQKV, "0"}, + {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, + {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "0"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + RunAndVerifyAttentionRoute("SdpaKernel=FLASH_ATTENTION", [&]() { + InvokePackedMultiHeadAttentionTest(true, true); + }); } -#endif if (kernel_type == AttentionKernelType::AttentionKernel_Unfused) { + if (!HasCudaEnvironment(530)) { + GTEST_SKIP() << "PackedMultiHeadAttention MATH route requires a CUDA device with FP16 support."; + } + ScopedEnvironmentVariables scoped_env_vars{ EnvVarMap{ {onnxruntime::contrib::attention::kDisableFlashAttention, "1"}, {onnxruntime::contrib::attention::kDisableTrtFlashAttention, "1"}, {onnxruntime::contrib::attention::kDisableFusedSelfAttention, "1"}, {onnxruntime::contrib::attention::kDisableFusedCrossAttention, "1"}, - {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}}}; - InvokePackedMultiHeadAttentionTest(true, true); - InvokePackedMultiHeadAttentionTest(false, false); + {onnxruntime::contrib::attention::kDisableMemoryEfficientAttention, "1"}, + {onnxruntime::contrib::attention::kEnableAttentionKernelDebugInfo, "1"}}}; + RunAndVerifyAttentionRoute("SdpaKernel=MATH", [&]() { + InvokePackedMultiHeadAttentionTest(true, true); + }); + RunAndVerifyAttentionRoute("SdpaKernel=MATH", [&]() { + InvokePackedMultiHeadAttentionTest(false, false); + }); } if (kernel_type == AttentionKernelType::AttentionKernel_Default) { @@ -310,7 +477,7 @@ TEST(PackedMultiHeadAttentionTest, Q_K_V_NoPadding_NoBias_trt) { AttentionKernelType::AttentionKernel_TrtFusedAttention); } -TEST(PackedMultiHeadAttentionTest, Q_K_V_NoPadding_Bias_AttnBias_cutlass) { +TEST(PackedMultiHeadAttentionTest, Q_K_V_NoPadding_Bias_AttnBias_InvalidHeadFallback) { AttentionTestData data; GetAttentionDataCutlassAttnBias(data); std::vector token_offset{0, 1, 2, 3, 4, 5, 6, 7}; @@ -404,31 +571,27 @@ TEST(PackedMultiHeadAttentionTest, PackedQKV_Padding_NoBias_cutlass) { AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention); } -#if USE_FLASH_ATTENTION TEST(PackedMultiHeadAttentionTest, PackedQKV_Padding_NoBias_FlashAttention) { - if (HasCudaEnvironment(800)) { - PackedAttentionTestData data; - GetPackedMultiHeadAttentionData_Batch2_HeadSize32_NoAttnBias(data); - std::vector empty_data = {}; - - RunPackedMultiHeadAttentionTest( - data.qkv_data, - empty_data, - empty_data, - empty_data, - data.token_offset, - data.cumulative_sequence_length, - data.fp16_output_data, - data.batch_size, - data.sequence_length, - data.hidden_size, - data.v_hidden_size, - data.num_heads, - data.token_count, - AttentionKernelType::AttentionKernel_FlashAttention); - } + PackedAttentionTestData data; + GetPackedMultiHeadAttentionData_Batch2_HeadSize32_NoAttnBias(data); + std::vector empty_data = {}; + + RunPackedMultiHeadAttentionTest( + data.qkv_data, + empty_data, + empty_data, + empty_data, + data.token_offset, + data.cumulative_sequence_length, + data.fp16_output_data, + data.batch_size, + data.sequence_length, + data.hidden_size, + data.v_hidden_size, + data.num_heads, + data.token_count, + AttentionKernelType::AttentionKernel_FlashAttention); } -#endif TEST(PackedMultiHeadAttentionTest, PackedQKV_Padding_NoBias_unfused) { PackedAttentionTestData data; diff --git a/onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_header_test.cc b/onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_header_test.cc new file mode 100644 index 0000000000000..c580f82758fa0 --- /dev/null +++ b/onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_header_test.cc @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Keep this header first and this translation unit free of ORT and test-framework +// headers. Both the in-tree CUDA target and the plugin-internal target compile it. +#include "contrib_ops/cuda/bert/packed_attention_workspace.h" + +#include + +namespace onnxruntime { +namespace test { + +using contrib::cuda::GetPackedAttentionQkvMaterializationIndexWidth; +using contrib::cuda::PackedAttentionProblem; +using contrib::cuda::PackedAttentionQkvMaterializationIndexWidth; +using contrib::cuda::PackedAttentionWorkspaceRecipe; +using contrib::cuda::PackedMultiHeadAttentionProblem; + +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_trivially_copyable_v); +static_assert(noexcept(GetPackedAttentionQkvMaterializationIndexWidth(4, 4))); +static_assert(GetPackedAttentionQkvMaterializationIndexWidth(4, 4) == + PackedAttentionQkvMaterializationIndexWidth::Vector4); +static_assert(GetPackedAttentionQkvMaterializationIndexWidth(2, 2) == + PackedAttentionQkvMaterializationIndexWidth::Vector2); +static_assert(GetPackedAttentionQkvMaterializationIndexWidth(1, 1) == + PackedAttentionQkvMaterializationIndexWidth::Scalar); + +void CompilePackedAttentionWorkspaceHeaderInIsolation() { + PackedAttentionProblem packed_attention; + PackedMultiHeadAttentionProblem packed_mha; + PackedAttentionWorkspaceRecipe recipe; + (void)packed_attention; + (void)packed_mha; + (void)recipe; +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_test.cc b/onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_test.cc new file mode 100644 index 0000000000000..5dbba47fef709 --- /dev/null +++ b/onnxruntime/test/providers/cuda/test_cases/packed_attention_workspace_test.cc @@ -0,0 +1,914 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#include +#include +#include +#include + +#include "contrib_ops/cuda/bert/packed_attention_workspace.h" + +namespace onnxruntime { +namespace test { + +using contrib::cuda::BuildPackedAttentionProblem; +using contrib::cuda::BuildPackedMultiHeadAttentionProblem; +using contrib::cuda::CheckedPackedAttentionAdd; +using contrib::cuda::CheckedPackedAttentionAlign; +using contrib::cuda::CheckedPackedAttentionMultiply; +using contrib::cuda::GetPackedAttentionWorkspaceRecipe; +using contrib::cuda::GetPackedMultiHeadAttentionWorkspaceRecipe; +using contrib::cuda::PackedAttentionBackend; +using contrib::cuda::PackedAttentionInputShapes; +using contrib::cuda::PackedAttentionProblem; +using contrib::cuda::PackedAttentionQkvMaterializationIndexWidth; +using contrib::cuda::PackedAttentionQkvWorkspaceLayout; +using contrib::cuda::PackedAttentionShape; +using contrib::cuda::PackedAttentionWorkspaceError; +using contrib::cuda::PackedMultiHeadAttentionInputShapes; +using contrib::cuda::PackedMultiHeadAttentionProblem; +using contrib::cuda::PackedMultiHeadAttentionQkvFormat; +using contrib::cuda::ValidatePackedAttentionWorkspaceRecipe; + +static_assert(std::is_trivially_copyable_v); +static_assert(std::is_trivially_copyable_v); + +namespace { + +PackedAttentionShape Shape(std::initializer_list dimensions) { + PackedAttentionShape shape; + shape.rank = dimensions.size(); + size_t index = 0; + for (int64_t dimension : dimensions) { + if (index < shape.dimensions.size()) { + shape.dimensions[index] = dimension; + } + ++index; + } + + return shape; +} + +PackedAttentionInputShapes ValidPackedAttentionInputs(int64_t token_count = 6) { + PackedAttentionInputShapes inputs; + inputs.input = Shape({token_count, 6}); + inputs.weights = Shape({6, 20}); + inputs.bias = Shape({20}); + inputs.token_offset = Shape({2, 4}); + inputs.cumulative_sequence_length = Shape({3}); + inputs.element_size = 2; + inputs.num_heads = 2; + inputs.qkv_hidden_sizes_count = 3; + inputs.qkv_hidden_sizes = {8, 8, 4}; + return inputs; +} + +PackedAttentionInputShapes ValidPackedAttentionEqualHeadsInputs(int64_t token_count = 6) { + auto inputs = ValidPackedAttentionInputs(token_count); + inputs.weights = Shape({6, 24}); + inputs.bias = Shape({24}); + inputs.qkv_hidden_sizes = {8, 8, 8}; + return inputs; +} + +PackedMultiHeadAttentionInputShapes ValidSeparateQkvInputs(int64_t token_count = 6) { + PackedMultiHeadAttentionInputShapes inputs; + inputs.query = Shape({token_count, 8}); + inputs.key = Shape({token_count, 8}); + inputs.value = Shape({token_count, 4}); + inputs.token_offset = Shape({2, 4}); + inputs.cumulative_sequence_length = Shape({3}); + inputs.element_size = 2; + inputs.num_heads = 2; + inputs.has_key = true; + inputs.has_value = true; + return inputs; +} + +PackedMultiHeadAttentionInputShapes ValidSeparateEqualQkvInputs(int64_t token_count = 6) { + auto inputs = ValidSeparateQkvInputs(token_count); + inputs.value = Shape({token_count, 8}); + return inputs; +} + +PackedMultiHeadAttentionInputShapes ValidPackedQkvInputs(int64_t token_count = 6) { + PackedMultiHeadAttentionInputShapes inputs; + inputs.query = Shape({token_count, 2, 3, 4}); + inputs.token_offset = Shape({2, 4}); + inputs.cumulative_sequence_length = Shape({3}); + inputs.element_size = 2; + inputs.num_heads = 2; + return inputs; +} + +} // namespace + +TEST(PackedAttentionWorkspaceTest, CheckedArithmeticHandlesZeroBoundaryAndOverflow) { + size_t result = 123; + EXPECT_TRUE(CheckedPackedAttentionAdd(0, 0, result).IsOK()); + EXPECT_EQ(result, 0U); + EXPECT_TRUE(CheckedPackedAttentionAdd(7, 9, result).IsOK()); + EXPECT_EQ(result, 16U); + EXPECT_EQ(CheckedPackedAttentionAdd(std::numeric_limits::max(), 1, result).error, + PackedAttentionWorkspaceError::Overflow); + + EXPECT_TRUE(CheckedPackedAttentionMultiply(0, std::numeric_limits::max(), result).IsOK()); + EXPECT_EQ(result, 0U); + EXPECT_TRUE(CheckedPackedAttentionMultiply(7, 9, result).IsOK()); + EXPECT_EQ(result, 63U); + EXPECT_EQ(CheckedPackedAttentionMultiply(std::numeric_limits::max(), 2, result).error, + PackedAttentionWorkspaceError::Overflow); + + EXPECT_TRUE(CheckedPackedAttentionAlign(0, 256, result).IsOK()); + EXPECT_EQ(result, 0U); + EXPECT_TRUE(CheckedPackedAttentionAlign(257, 256, result).IsOK()); + EXPECT_EQ(result, 512U); + EXPECT_TRUE(CheckedPackedAttentionAlign(std::numeric_limits::max() - 255, 256, result).IsOK()); + EXPECT_EQ(result, std::numeric_limits::max() - 255); + EXPECT_EQ(CheckedPackedAttentionAlign(std::numeric_limits::max(), 256, result).error, + PackedAttentionWorkspaceError::Overflow); + EXPECT_EQ(CheckedPackedAttentionAlign(1, 0, result).error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +enum class LegacyParityOperator { + PackedAttention, + PackedMultiHeadAttention, +}; + +struct LegacyParityCase { + const char* name; + LegacyParityOperator op; + PackedAttentionBackend backend; + int64_t token_count; + bool expected_no_qkv_workspace; + size_t expected_projection_bytes; + size_t expected_qkv_capacity_bytes; + size_t expected_planar_q_bytes; + size_t expected_interleaved_qkv_bytes; + size_t expected_backend_offset_bytes; + size_t expected_backend_bytes; + size_t expected_attention_workspace_bytes; + PackedAttentionQkvWorkspaceLayout expected_layout; +}; + +class PackedAttentionLegacyParityTest : public testing::TestWithParam {}; + +TEST_P(PackedAttentionLegacyParityTest, MatchesHandCalculatedLegacyComponents) { + const LegacyParityCase& test_case = GetParam(); + contrib::cuda::PackedAttentionWorkspaceResult workspace_result; + + if (test_case.op == LegacyParityOperator::PackedAttention) { + auto problem_result = BuildPackedAttentionProblem( + ValidPackedAttentionEqualHeadsInputs(test_case.token_count)); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + problem_result.problem.backend = test_case.backend; + problem_result.problem.trt_runner_available = test_case.backend == PackedAttentionBackend::Trt; + workspace_result = GetPackedAttentionWorkspaceRecipe(problem_result.problem); + } else { + PackedMultiHeadAttentionInputShapes inputs; + if (test_case.backend == PackedAttentionBackend::Trt) { + inputs = test_case.expected_no_qkv_workspace + ? ValidPackedQkvInputs(test_case.token_count) + : ValidSeparateEqualQkvInputs(test_case.token_count); + } else if (test_case.backend == PackedAttentionBackend::Flash || + test_case.backend == PackedAttentionBackend::MemoryEfficient) { + inputs = test_case.expected_no_qkv_workspace + ? ValidSeparateEqualQkvInputs(test_case.token_count) + : ValidPackedQkvInputs(test_case.token_count); + } else { + inputs = ValidPackedQkvInputs(test_case.token_count); + } + + auto problem_result = BuildPackedMultiHeadAttentionProblem(inputs); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + problem_result.problem.backend = test_case.backend; + problem_result.problem.trt_runner_available = test_case.backend == PackedAttentionBackend::Trt; + workspace_result = GetPackedMultiHeadAttentionWorkspaceRecipe(problem_result.problem); + } + + ASSERT_TRUE(workspace_result.status.IsOK()) << workspace_result.status.message; + const auto& recipe = workspace_result.recipe; + EXPECT_EQ(recipe.no_qkv_workspace, test_case.expected_no_qkv_workspace); + EXPECT_EQ(recipe.projection_bytes, test_case.expected_projection_bytes); + EXPECT_EQ(recipe.qkv_capacity_bytes, test_case.expected_qkv_capacity_bytes); + EXPECT_EQ(recipe.q_bytes, test_case.expected_planar_q_bytes); + EXPECT_EQ(recipe.interleaved_qkv_bytes, test_case.expected_interleaved_qkv_bytes); + EXPECT_EQ(recipe.backend_workspace_offset_bytes, test_case.expected_backend_offset_bytes); + EXPECT_EQ(recipe.backend_workspace_bytes, test_case.expected_backend_bytes); + EXPECT_EQ(recipe.attention_workspace_bytes, test_case.expected_attention_workspace_bytes); + EXPECT_EQ(recipe.qkv_layout, test_case.expected_layout); + EXPECT_TRUE(ValidatePackedAttentionWorkspaceRecipe(recipe).IsOK()); +} + +INSTANTIATE_TEST_SUITE_P( + HandCalculatedMatrix, + PackedAttentionLegacyParityTest, + testing::Values( + // B=2, S=4, N=2, H=Hv=4, sizeof(T)=2: + // QKV capacity=384, Flash LSE=64, unfused scratch=align256(128)=256. + LegacyParityCase{"PmhaFlashDirectT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Flash, 6, true, 0, 0, 0, 0, 0, 64, 64, + PackedAttentionQkvWorkspaceLayout::None}, + LegacyParityCase{"PmhaFlashPackedT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Flash, 6, false, 0, 384, 96, 0, 288, 64, 448, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PmhaMeaDirectT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::MemoryEfficient, 6, true, 0, 0, 0, 0, 0, 0, 0, + PackedAttentionQkvWorkspaceLayout::None}, + LegacyParityCase{"PmhaMeaPackedT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::MemoryEfficient, 6, false, 0, 384, 96, 0, 288, 0, 384, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PmhaTrtDirectT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Trt, 6, true, 0, 0, 0, 0, 0, 0, 0, + PackedAttentionQkvWorkspaceLayout::None}, + LegacyParityCase{"PmhaTrtMaterializedT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Trt, 6, false, 0, 384, 0, 288, 0, 0, 384, + PackedAttentionQkvWorkspaceLayout::InterleavedTn3h}, + LegacyParityCase{"PmhaUnfusedT6", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Unfused, 6, false, 0, 384, 128, 0, 384, 256, 896, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PmhaFlashDirectT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Flash, 8, true, 0, 0, 0, 0, 0, 64, 64, + PackedAttentionQkvWorkspaceLayout::None}, + LegacyParityCase{"PmhaFlashPackedT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Flash, 8, false, 0, 384, 128, 0, 384, 64, 448, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PmhaMeaDirectT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::MemoryEfficient, 8, true, 0, 0, 0, 0, 0, 0, 0, + PackedAttentionQkvWorkspaceLayout::None}, + LegacyParityCase{"PmhaMeaPackedT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::MemoryEfficient, 8, false, 0, 384, 128, 0, 384, 0, 384, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PmhaTrtDirectT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Trt, 8, true, 0, 0, 0, 0, 0, 0, 0, + PackedAttentionQkvWorkspaceLayout::None}, + LegacyParityCase{"PmhaTrtMaterializedT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Trt, 8, false, 0, 384, 0, 384, 0, 0, 384, + PackedAttentionQkvWorkspaceLayout::InterleavedTn3h}, + LegacyParityCase{"PmhaUnfusedT8", LegacyParityOperator::PackedMultiHeadAttention, + PackedAttentionBackend::Unfused, 8, false, 0, 384, 128, 0, 384, 256, 896, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PaMeaT6", LegacyParityOperator::PackedAttention, + PackedAttentionBackend::MemoryEfficient, 6, false, 288, 384, 96, 0, 288, 0, 384, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PaTrtT6", LegacyParityOperator::PackedAttention, + PackedAttentionBackend::Trt, 6, false, 288, 384, 0, 288, 0, 0, 384, + PackedAttentionQkvWorkspaceLayout::InterleavedTn3h}, + LegacyParityCase{"PaUnfusedT6", LegacyParityOperator::PackedAttention, + PackedAttentionBackend::Unfused, 6, false, 288, 384, 128, 0, 384, 256, 896, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PaMeaT8", LegacyParityOperator::PackedAttention, + PackedAttentionBackend::MemoryEfficient, 8, false, 384, 384, 128, 0, 384, 0, 384, + PackedAttentionQkvWorkspaceLayout::Planar}, + LegacyParityCase{"PaTrtT8", LegacyParityOperator::PackedAttention, + PackedAttentionBackend::Trt, 8, false, 384, 384, 0, 384, 0, 0, 384, + PackedAttentionQkvWorkspaceLayout::InterleavedTn3h}, + LegacyParityCase{"PaUnfusedT8", LegacyParityOperator::PackedAttention, + PackedAttentionBackend::Unfused, 8, false, 384, 384, 128, 0, 384, 256, 896, + PackedAttentionQkvWorkspaceLayout::Planar}), + [](const testing::TestParamInfo& info) { + return std::string(info.param.name); + }); + +TEST(PackedAttentionWorkspaceTest, UnfusedGoldenPreservesLegacyCapacityWhenPacked) { + auto problem_result = BuildPackedAttentionProblem(ValidPackedAttentionInputs()); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + problem_result.problem.backend = PackedAttentionBackend::Unfused; + + auto workspace_result = GetPackedAttentionWorkspaceRecipe(problem_result.problem); + ASSERT_TRUE(workspace_result.status.IsOK()) << workspace_result.status.message; + const auto& recipe = workspace_result.recipe; + + // Projection: T * (Q + K + V) * 2 = 6 * 20 * 2 = 240. + EXPECT_EQ(recipe.projection_bytes, 240U); + EXPECT_EQ(recipe.projection_m, 6); + EXPECT_EQ(recipe.projection_n, 20); + EXPECT_EQ(recipe.projection_k, 6); + + // QKV capacity: B*S*N*(H+H+Hv)*2 = 2*4*2*10*2 = 320. + // Each attention scratch: align256(2*2*2*4*4) = 256. Total = 320 + 2*256 = 832. + EXPECT_EQ(recipe.qkv_capacity_bytes, 320U); + EXPECT_EQ(recipe.qkv_layout, PackedAttentionQkvWorkspaceLayout::Planar); + EXPECT_EQ(recipe.q_offset_bytes, 0U); + EXPECT_EQ(recipe.q_bytes, 128U); + EXPECT_EQ(recipe.k_offset_bytes, 128U); + EXPECT_EQ(recipe.v_offset_bytes, 256U); + EXPECT_EQ(recipe.v_bytes, 64U); + EXPECT_EQ(recipe.backend_workspace_offset_bytes, 320U); + EXPECT_EQ(recipe.backend_workspace_bytes, 256U); + EXPECT_TRUE(recipe.has_second_scratch); + EXPECT_EQ(recipe.second_scratch_offset_bytes, 576U); + EXPECT_EQ(recipe.attention_workspace_bytes, 832U); +} + +TEST(PackedAttentionWorkspaceTest, LegacyAttentionComponentMatchesForTEqualPaddedCapacity) { + auto pa_problem = BuildPackedAttentionProblem(ValidPackedAttentionInputs(8)); + ASSERT_TRUE(pa_problem.status.IsOK()) << pa_problem.status.message; + pa_problem.problem.backend = PackedAttentionBackend::Unfused; + auto pa_workspace = GetPackedAttentionWorkspaceRecipe(pa_problem.problem); + ASSERT_TRUE(pa_workspace.status.IsOK()) << pa_workspace.status.message; + + auto pmha_problem = BuildPackedMultiHeadAttentionProblem(ValidSeparateQkvInputs(8)); + ASSERT_TRUE(pmha_problem.status.IsOK()) << pmha_problem.status.message; + pmha_problem.problem.backend = PackedAttentionBackend::Unfused; + auto pmha_workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(pmha_problem.problem); + ASSERT_TRUE(pmha_workspace.status.IsOK()) << pmha_workspace.status.message; + + // Both operators retain the shared legacy B*S attention allocation. + EXPECT_EQ(pa_workspace.recipe.attention_workspace_bytes, 832U); + EXPECT_EQ(pmha_workspace.recipe.attention_workspace_bytes, 832U); + EXPECT_EQ(pa_workspace.recipe.projection_bytes, 320U); + EXPECT_EQ(pmha_workspace.recipe.projection_bytes, 0U); +} + +TEST(PackedAttentionWorkspaceTest, TokenCountGreaterThanPaddedCapacityIsRejected) { + auto pa_result = BuildPackedAttentionProblem(ValidPackedAttentionInputs(9)); + EXPECT_EQ(pa_result.status.error, PackedAttentionWorkspaceError::InvalidArgument); + + auto pmha_result = BuildPackedMultiHeadAttentionProblem(ValidSeparateQkvInputs(9)); + EXPECT_EQ(pmha_result.status.error, PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, PackedMetadataShapesMustMatchBatchAndSequence) { + auto pa_inputs = ValidPackedAttentionInputs(); + pa_inputs.token_offset = Shape({8}); + EXPECT_EQ(BuildPackedAttentionProblem(pa_inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + pa_inputs = ValidPackedAttentionInputs(); + pa_inputs.cumulative_sequence_length = Shape({4}); + EXPECT_EQ(BuildPackedAttentionProblem(pa_inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + auto pmha_inputs = ValidSeparateQkvInputs(); + pmha_inputs.cumulative_sequence_length = Shape({2}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(pmha_inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, LargeAttentionBiasDimensionIsNotReportedAsRankError) { + auto inputs = ValidPackedAttentionInputs(); + inputs.has_attention_bias = true; + inputs.attention_bias = + Shape({1, 2, static_cast(std::numeric_limits::max()) + 1, 4}); + + const auto result = BuildPackedAttentionProblem(inputs); + EXPECT_EQ(result.status.error, PackedAttentionWorkspaceError::InvalidArgument); + EXPECT_NE(std::string(result.status.message).find("int32 CUDA ABI"), std::string::npos) + << result.status.message; + EXPECT_EQ(std::string(result.status.message).find("rank"), std::string::npos) + << result.status.message; +} + +TEST(PackedAttentionWorkspaceTest, FusedViewsUseTButTotalRetainsBatchSequenceCapacity) { + auto problem_result = BuildPackedMultiHeadAttentionProblem(ValidPackedQkvInputs()); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + problem_result.problem.backend = PackedAttentionBackend::Flash; + + auto workspace_result = GetPackedMultiHeadAttentionWorkspaceRecipe(problem_result.problem); + ASSERT_TRUE(workspace_result.status.IsOK()) << workspace_result.status.message; + const auto& recipe = workspace_result.recipe; + + // Legacy QKV capacity uses B*S: 2*4*2*(4+4+4)*2 = 384. + EXPECT_EQ(recipe.qkv_capacity_bytes, 384U); + // Runtime views use T: each Q/K/V view is 6*2*4*2 = 96. + EXPECT_EQ(recipe.qkv_layout, PackedAttentionQkvWorkspaceLayout::Planar); + EXPECT_EQ(recipe.q_bytes, 96U); + EXPECT_EQ(recipe.k_offset_bytes, 96U); + EXPECT_EQ(recipe.v_offset_bytes, 192U); + EXPECT_EQ(recipe.backend_workspace_offset_bytes, 288U); + // Flash LSE is 4*B*S*N = 64, while total remains capacity + LSE. + EXPECT_EQ(recipe.backend_workspace_bytes, 64U); + EXPECT_EQ(recipe.attention_workspace_bytes, 448U); +} + +TEST(PackedAttentionWorkspaceTest, ProjectionUsesActualGemmDimensionsAndPmhaHasNone) { + auto pa_problem = BuildPackedAttentionProblem(ValidPackedAttentionInputs()); + ASSERT_TRUE(pa_problem.status.IsOK()) << pa_problem.status.message; + auto pa_workspace = GetPackedAttentionWorkspaceRecipe(pa_problem.problem); + ASSERT_TRUE(pa_workspace.status.IsOK()) << pa_workspace.status.message; + + // input_hidden=6, Q hidden=8, V hidden=4. + EXPECT_EQ(pa_workspace.recipe.projection_k, 6); + EXPECT_EQ(pa_workspace.recipe.projection_n, 20); + EXPECT_EQ(pa_workspace.recipe.projection_bytes, 240U); + + auto pmha_problem = BuildPackedMultiHeadAttentionProblem(ValidSeparateQkvInputs()); + ASSERT_TRUE(pmha_problem.status.IsOK()) << pmha_problem.status.message; + auto pmha_workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(pmha_problem.problem); + ASSERT_TRUE(pmha_workspace.status.IsOK()) << pmha_workspace.status.message; + EXPECT_EQ(pmha_problem.problem.hidden_size, 8); + EXPECT_EQ(pmha_problem.problem.v_hidden_size, 4); + EXPECT_EQ(pmha_workspace.recipe.projection_bytes, 0U); +} + +TEST(PackedAttentionWorkspaceTest, PmhaDirectQkvRoutesHaveZeroProjectionAndNoQkvWorkspace) { + auto packed_problem = BuildPackedMultiHeadAttentionProblem(ValidPackedQkvInputs()); + ASSERT_TRUE(packed_problem.status.IsOK()) << packed_problem.status.message; + packed_problem.problem.backend = PackedAttentionBackend::Trt; + packed_problem.problem.trt_runner_available = true; + auto trt_workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(packed_problem.problem); + ASSERT_TRUE(trt_workspace.status.IsOK()) << trt_workspace.status.message; + EXPECT_TRUE(trt_workspace.recipe.no_qkv_workspace); + EXPECT_EQ(trt_workspace.recipe.qkv_layout, PackedAttentionQkvWorkspaceLayout::None); + EXPECT_EQ(trt_workspace.recipe.projection_bytes, 0U); + EXPECT_EQ(trt_workspace.recipe.attention_workspace_bytes, 0U); + + auto separate_problem = BuildPackedMultiHeadAttentionProblem(ValidSeparateQkvInputs()); + ASSERT_TRUE(separate_problem.status.IsOK()) << separate_problem.status.message; + separate_problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + auto mea_workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(separate_problem.problem); + ASSERT_TRUE(mea_workspace.status.IsOK()) << mea_workspace.status.message; + EXPECT_TRUE(mea_workspace.recipe.no_qkv_workspace); + EXPECT_EQ(mea_workspace.recipe.qkv_layout, PackedAttentionQkvWorkspaceLayout::None); + EXPECT_EQ(mea_workspace.recipe.attention_workspace_bytes, 0U); +} + +TEST(PackedAttentionWorkspaceTest, DirectQkvFusedRoutesDoNotRequireInt32SequenceSquare) { + constexpr int64_t kSequenceLength = 46341; // S*S is greater than INT32_MAX. + + auto separate_inputs = ValidSeparateEqualQkvInputs(1); + separate_inputs.token_offset = Shape({1, kSequenceLength}); + separate_inputs.cumulative_sequence_length = Shape({2}); + auto separate_problem = BuildPackedMultiHeadAttentionProblem(separate_inputs); + ASSERT_TRUE(separate_problem.status.IsOK()) << separate_problem.status.message; + + for (PackedAttentionBackend backend : + {PackedAttentionBackend::Flash, PackedAttentionBackend::MemoryEfficient}) { + separate_problem.problem.backend = backend; + auto workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(separate_problem.problem); + ASSERT_TRUE(workspace.status.IsOK()) << workspace.status.message; + EXPECT_TRUE(workspace.recipe.no_qkv_workspace); + } + + auto packed_inputs = ValidPackedQkvInputs(1); + packed_inputs.token_offset = Shape({1, kSequenceLength}); + packed_inputs.cumulative_sequence_length = Shape({2}); + auto packed_problem = BuildPackedMultiHeadAttentionProblem(packed_inputs); + ASSERT_TRUE(packed_problem.status.IsOK()) << packed_problem.status.message; + packed_problem.problem.backend = PackedAttentionBackend::Trt; + packed_problem.problem.trt_runner_available = true; + auto trt_workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(packed_problem.problem); + ASSERT_TRUE(trt_workspace.status.IsOK()) << trt_workspace.status.message; + EXPECT_TRUE(trt_workspace.recipe.no_qkv_workspace); + + separate_problem.problem.backend = PackedAttentionBackend::Unfused; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(separate_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, MeaAttentionBiasStridesUseInt64) { + constexpr int64_t kSequenceLength = 32768; + constexpr int64_t kNumHeads = 2; + static_assert(kNumHeads * kSequenceLength * kSequenceLength > + std::numeric_limits::max()); + + auto inputs = ValidPackedQkvInputs(1); + inputs.token_offset = Shape({1, kSequenceLength}); + inputs.cumulative_sequence_length = Shape({2}); + inputs.attention_bias = Shape({1, kNumHeads, kSequenceLength, kSequenceLength}); + inputs.has_attention_bias = true; + + auto problem = BuildPackedMultiHeadAttentionProblem(inputs); + ASSERT_TRUE(problem.status.IsOK()) << problem.status.message; + problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + auto workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(problem.problem); + ASSERT_TRUE(workspace.status.IsOK()) << workspace.status.message; + EXPECT_EQ(workspace.recipe.qkv_layout, PackedAttentionQkvWorkspaceLayout::Planar); +} + +TEST(PackedAttentionWorkspaceTest, TrtMaterializationExposesOnlyInterleavedTn3hRegion) { + auto problem = BuildPackedMultiHeadAttentionProblem(ValidSeparateEqualQkvInputs()); + ASSERT_TRUE(problem.status.IsOK()) << problem.status.message; + problem.problem.backend = PackedAttentionBackend::Trt; + problem.problem.trt_runner_available = true; + + auto workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(problem.problem); + ASSERT_TRUE(workspace.status.IsOK()) << workspace.status.message; + const auto& recipe = workspace.recipe; + + EXPECT_EQ(recipe.qkv_layout, PackedAttentionQkvWorkspaceLayout::InterleavedTn3h); + EXPECT_EQ(recipe.interleaved_qkv_offset_bytes, 0U); + // Producer address: (((t * N + n) * 3 + component) * H + h) * sizeof(fp16). + // With T=6, N=2, H=4, the end of V for the last head is byte 288. + constexpr size_t kProducerRegionEnd = ((((6U - 1) * 2 + (2U - 1)) * 3 + 2) * 4 + 4) * 2; + EXPECT_EQ(recipe.interleaved_qkv_bytes, kProducerRegionEnd); + EXPECT_EQ(recipe.q_offset_bytes, 0U); + EXPECT_EQ(recipe.q_bytes, 0U); + EXPECT_EQ(recipe.k_offset_bytes, 0U); + EXPECT_EQ(recipe.k_bytes, 0U); + EXPECT_EQ(recipe.v_offset_bytes, 0U); + EXPECT_EQ(recipe.v_bytes, 0U); +} + +TEST(PackedAttentionWorkspaceTest, MemoryEfficientAccumulatorUsesLegacyCapacity) { + auto inputs = ValidSeparateQkvInputs(); + inputs.value = Shape({6, 320}); + inputs.has_bias = true; + inputs.bias = Shape({336}); + auto problem_result = BuildPackedMultiHeadAttentionProblem(inputs); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + problem_result.problem.backend = PackedAttentionBackend::MemoryEfficient; + + auto workspace_result = GetPackedMultiHeadAttentionWorkspaceRecipe(problem_result.problem); + ASSERT_TRUE(workspace_result.status.IsOK()) << workspace_result.status.message; + const auto& recipe = workspace_result.recipe; + + // QKV capacity: 2*4*2*(4+4+160)*2 = 5376. T-view ends at 6*2*168*2 = 4032. + // FP32 accumulator: 4*2*4*2*160 = 10240. Legacy total = 5376 + 10240 = 15616. + EXPECT_EQ(recipe.qkv_capacity_bytes, 5376U); + EXPECT_EQ(recipe.backend_workspace_offset_bytes, 4032U); + EXPECT_EQ(recipe.backend_workspace_bytes, 10240U); + EXPECT_EQ(recipe.attention_workspace_bytes, 15616U); +} + +TEST(PackedAttentionWorkspaceTest, PackedQkvAxesAreStrictlyValidated) { + auto inputs = ValidPackedQkvInputs(); + inputs.query = Shape({6, 1, 3, 4}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs = ValidPackedQkvInputs(); + inputs.query = Shape({6, 2, 2, 4}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs = ValidPackedQkvInputs(); + inputs.query = Shape({6, 2, 3}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, HeadDivisibilityAndSeparateQkvConsistencyAreValidated) { + auto inputs = ValidSeparateQkvInputs(); + inputs.query = Shape({6, 7}); + inputs.key = Shape({6, 7}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs = ValidSeparateQkvInputs(); + inputs.key = Shape({6, 4}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs = ValidSeparateQkvInputs(); + inputs.value = Shape({5, 4}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs = ValidSeparateQkvInputs(); + inputs.has_value = false; + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + auto pa_inputs = ValidPackedAttentionInputs(); + pa_inputs.qkv_hidden_sizes = {7, 7, 6}; + EXPECT_EQ(BuildPackedAttentionProblem(pa_inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, Int32NarrowingAndDerivedProductsAreValidated) { + auto inputs = ValidPackedQkvInputs(); + inputs.query = Shape({static_cast(std::numeric_limits::max()) + 1, 2, 3, 4}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs = ValidPackedQkvInputs(0); + inputs.query = Shape({0, 50000, 3, 1}); + inputs.num_heads = 50000; + inputs.token_offset = Shape({1, 50000}); + inputs.cumulative_sequence_length = Shape({2}); + auto large_ns_problem = BuildPackedMultiHeadAttentionProblem(inputs); + ASSERT_TRUE(large_ns_problem.status.IsOK()) << large_ns_problem.status.message; + large_ns_problem.problem.backend = PackedAttentionBackend::Unfused; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(large_ns_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + inputs.query = Shape({0, 50000, 3, 50000}); + EXPECT_EQ(BuildPackedMultiHeadAttentionProblem(inputs).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + auto large_bs_inputs = ValidSeparateEqualQkvInputs(0); + large_bs_inputs.query = Shape({0, 8}); + large_bs_inputs.key = Shape({0, 8}); + large_bs_inputs.value = Shape({0, 8}); + large_bs_inputs.token_offset = Shape({50000, 50000}); + large_bs_inputs.cumulative_sequence_length = Shape({50001}); + auto large_bs_problem = BuildPackedMultiHeadAttentionProblem(large_bs_inputs); + ASSERT_TRUE(large_bs_problem.status.IsOK()) << large_bs_problem.status.message; + large_bs_problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + EXPECT_TRUE(GetPackedMultiHeadAttentionWorkspaceRecipe(large_bs_problem.problem).status.IsOK()); + large_bs_problem.problem.backend = PackedAttentionBackend::Unfused; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(large_bs_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + constexpr int64_t kLargeTokenCount = 100000; + constexpr int64_t kLargeNumHeads = 4000; + auto materialized_inputs = ValidPackedQkvInputs(kLargeTokenCount); + materialized_inputs.query = Shape({kLargeTokenCount, kLargeNumHeads, 3, 8}); + materialized_inputs.num_heads = kLargeNumHeads; + materialized_inputs.token_offset = Shape({1, kLargeTokenCount}); + materialized_inputs.cumulative_sequence_length = Shape({2}); + auto materialized_problem = BuildPackedMultiHeadAttentionProblem(materialized_inputs); + ASSERT_TRUE(materialized_problem.status.IsOK()) << materialized_problem.status.message; + materialized_problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(materialized_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + PackedMultiHeadAttentionInputShapes direct_inputs; + direct_inputs.query = Shape({kLargeTokenCount, kLargeNumHeads * 8}); + direct_inputs.key = direct_inputs.query; + direct_inputs.value = direct_inputs.query; + direct_inputs.token_offset = Shape({1, kLargeTokenCount}); + direct_inputs.cumulative_sequence_length = Shape({2}); + direct_inputs.element_size = 2; + direct_inputs.num_heads = kLargeNumHeads; + direct_inputs.has_key = true; + direct_inputs.has_value = true; + auto direct_problem = BuildPackedMultiHeadAttentionProblem(direct_inputs); + ASSERT_TRUE(direct_problem.status.IsOK()) << direct_problem.status.message; + direct_problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + auto direct_workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(direct_problem.problem); + ASSERT_TRUE(direct_workspace.status.IsOK()) << direct_workspace.status.message; + EXPECT_TRUE(direct_workspace.recipe.no_qkv_workspace); +} + +TEST(PackedAttentionWorkspaceTest, AllQkvMaterializationProducersUseSelectedVectorIndexWidth) { + auto vector4_inputs = ValidSeparateEqualQkvInputs(); + auto vector4_problem = BuildPackedMultiHeadAttentionProblem(vector4_inputs); + ASSERT_TRUE(vector4_problem.status.IsOK()) << vector4_problem.status.message; + EXPECT_EQ(vector4_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector4); + + auto vector2_inputs = ValidSeparateEqualQkvInputs(); + vector2_inputs.query = Shape({6, 12}); + vector2_inputs.key = Shape({6, 12}); + vector2_inputs.value = Shape({6, 12}); + auto vector2_problem = BuildPackedMultiHeadAttentionProblem(vector2_inputs); + ASSERT_TRUE(vector2_problem.status.IsOK()) << vector2_problem.status.message; + EXPECT_EQ(vector2_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector2); + + auto scalar_inputs = ValidSeparateEqualQkvInputs(); + scalar_inputs.query = Shape({6, 6}); + scalar_inputs.key = Shape({6, 6}); + scalar_inputs.value = Shape({6, 6}); + auto scalar_problem = BuildPackedMultiHeadAttentionProblem(scalar_inputs); + ASSERT_TRUE(scalar_problem.status.IsOK()) << scalar_problem.status.message; + EXPECT_EQ(scalar_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Scalar); + scalar_problem.problem.backend = PackedAttentionBackend::Trt; + scalar_problem.problem.trt_runner_available = true; + scalar_problem.problem.qkv_materialization_index_width = + PackedAttentionQkvMaterializationIndexWidth::Vector4; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(scalar_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + auto packed_problem = BuildPackedMultiHeadAttentionProblem(ValidPackedQkvInputs()); + ASSERT_TRUE(packed_problem.status.IsOK()) << packed_problem.status.message; + EXPECT_EQ(packed_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector4); + + auto packed_vector2_inputs = ValidPackedQkvInputs(); + packed_vector2_inputs.query = Shape({6, 2, 3, 2}); + auto packed_vector2_problem = BuildPackedMultiHeadAttentionProblem(packed_vector2_inputs); + ASSERT_TRUE(packed_vector2_problem.status.IsOK()) << packed_vector2_problem.status.message; + EXPECT_EQ(packed_vector2_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector2); + + auto packed_scalar_inputs = ValidPackedQkvInputs(); + packed_scalar_inputs.query = Shape({6, 2, 3, 1}); + auto packed_scalar_problem = BuildPackedMultiHeadAttentionProblem(packed_scalar_inputs); + ASSERT_TRUE(packed_scalar_problem.status.IsOK()) << packed_scalar_problem.status.message; + EXPECT_EQ(packed_scalar_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Scalar); + + auto pa_problem = BuildPackedAttentionProblem(ValidPackedAttentionEqualHeadsInputs()); + ASSERT_TRUE(pa_problem.status.IsOK()) << pa_problem.status.message; + EXPECT_EQ(pa_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector4); + + auto pa_vector2_inputs = ValidPackedAttentionEqualHeadsInputs(); + pa_vector2_inputs.weights = Shape({6, 12}); + pa_vector2_inputs.bias = Shape({12}); + pa_vector2_inputs.qkv_hidden_sizes = {4, 4, 4}; + auto pa_vector2_problem = BuildPackedAttentionProblem(pa_vector2_inputs); + ASSERT_TRUE(pa_vector2_problem.status.IsOK()) << pa_vector2_problem.status.message; + EXPECT_EQ(pa_vector2_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector2); + + auto pa_scalar_inputs = ValidPackedAttentionEqualHeadsInputs(); + pa_scalar_inputs.weights = Shape({6, 6}); + pa_scalar_inputs.bias = Shape({6}); + pa_scalar_inputs.qkv_hidden_sizes = {2, 2, 2}; + auto pa_scalar_problem = BuildPackedAttentionProblem(pa_scalar_inputs); + ASSERT_TRUE(pa_scalar_problem.status.IsOK()) << pa_scalar_problem.status.message; + EXPECT_EQ(pa_scalar_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Scalar); +} + +TEST(PackedAttentionWorkspaceTest, PackedFormatMaterializationBoundariesUseProducerIndexWidth) { + constexpr int64_t kNumHeads = 32768; + constexpr int64_t kLargestVectorSafeSequenceLength = 21845; + static_assert(kLargestVectorSafeSequenceLength * kNumHeads * 3 <= + std::numeric_limits::max()); + static_assert((kLargestVectorSafeSequenceLength + 1) * kNumHeads * 3 > + std::numeric_limits::max()); + static_assert((kLargestVectorSafeSequenceLength + 1) * + (kLargestVectorSafeSequenceLength + 1) <= + std::numeric_limits::max()); + + const auto make_inputs = [](int64_t sequence_length, int64_t head_size) { + PackedMultiHeadAttentionInputShapes inputs; + inputs.query = Shape({0, kNumHeads, 3, head_size}); + inputs.token_offset = Shape({1, sequence_length}); + inputs.cumulative_sequence_length = Shape({2}); + inputs.element_size = 2; + inputs.num_heads = kNumHeads; + return inputs; + }; + + struct BoundaryCase { + int64_t head_size; + PackedAttentionQkvMaterializationIndexWidth expected_width; + }; + + for (const auto& test_case : + {BoundaryCase{4, PackedAttentionQkvMaterializationIndexWidth::Vector4}, + BoundaryCase{2, PackedAttentionQkvMaterializationIndexWidth::Vector2}, + BoundaryCase{1, PackedAttentionQkvMaterializationIndexWidth::Scalar}}) { + SCOPED_TRACE(test_case.head_size); + auto boundary_problem = BuildPackedMultiHeadAttentionProblem( + make_inputs(kLargestVectorSafeSequenceLength, test_case.head_size)); + ASSERT_TRUE(boundary_problem.status.IsOK()) << boundary_problem.status.message; + ASSERT_EQ(boundary_problem.problem.qkv_materialization_index_width, + test_case.expected_width); + boundary_problem.problem.backend = PackedAttentionBackend::Unfused; + const auto boundary_workspace = + GetPackedMultiHeadAttentionWorkspaceRecipe(boundary_problem.problem); + ASSERT_TRUE(boundary_workspace.status.IsOK()) << boundary_workspace.status.message; + + auto beyond_problem = BuildPackedMultiHeadAttentionProblem( + make_inputs(kLargestVectorSafeSequenceLength + 1, test_case.head_size)); + ASSERT_TRUE(beyond_problem.status.IsOK()) << beyond_problem.status.message; + ASSERT_EQ(beyond_problem.problem.qkv_materialization_index_width, + test_case.expected_width); + beyond_problem.problem.backend = PackedAttentionBackend::Unfused; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(beyond_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + } +} + +TEST(PackedAttentionWorkspaceTest, PackedAttentionProjectionUsesVectorUnitMaterializationBoundary) { + constexpr int64_t kNumHeads = 32768; + constexpr int64_t kHeadSize = 4; + constexpr int64_t kLargestVectorSafeTokenCount = 21845; + constexpr int64_t kHiddenSize = kNumHeads * kHeadSize; + constexpr int64_t kProjectionSize = 3 * kHiddenSize; + static_assert(kLargestVectorSafeTokenCount * kNumHeads * 3 <= + std::numeric_limits::max()); + static_assert((kLargestVectorSafeTokenCount + 1) * kNumHeads * 3 > + std::numeric_limits::max()); + static_assert(kLargestVectorSafeTokenCount * kProjectionSize > + std::numeric_limits::max()); + + const auto make_inputs = [](int64_t token_count) { + PackedAttentionInputShapes inputs; + inputs.input = Shape({token_count, 1}); + inputs.weights = Shape({1, kProjectionSize}); + inputs.bias = Shape({kProjectionSize}); + inputs.token_offset = Shape({1, token_count}); + inputs.cumulative_sequence_length = Shape({2}); + inputs.element_size = 2; + inputs.num_heads = kNumHeads; + inputs.qkv_hidden_sizes_count = 3; + inputs.qkv_hidden_sizes = {kHiddenSize, kHiddenSize, kHiddenSize}; + return inputs; + }; + + auto boundary_problem = + BuildPackedAttentionProblem(make_inputs(kLargestVectorSafeTokenCount)); + ASSERT_TRUE(boundary_problem.status.IsOK()) << boundary_problem.status.message; + ASSERT_EQ(boundary_problem.problem.qkv_materialization_index_width, + PackedAttentionQkvMaterializationIndexWidth::Vector4); + boundary_problem.problem.backend = PackedAttentionBackend::Unfused; + const auto boundary_workspace = GetPackedAttentionWorkspaceRecipe(boundary_problem.problem); + ASSERT_TRUE(boundary_workspace.status.IsOK()) << boundary_workspace.status.message; + EXPECT_EQ(boundary_workspace.recipe.projection_m, kLargestVectorSafeTokenCount); + EXPECT_EQ(boundary_workspace.recipe.projection_n, kProjectionSize); + + auto beyond_problem = + BuildPackedAttentionProblem(make_inputs(kLargestVectorSafeTokenCount + 1)); + ASSERT_TRUE(beyond_problem.status.IsOK()) << beyond_problem.status.message; + beyond_problem.problem.backend = PackedAttentionBackend::Unfused; + EXPECT_EQ(GetPackedAttentionWorkspaceRecipe(beyond_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, MemoryEfficientQueryBlockRoundingIsValidatedAtInt32Boundary) { + constexpr int64_t kLargestRound64SafeSequenceLength = + static_cast(std::numeric_limits::max()) - 63; + + const auto make_problem = [](int64_t sequence_length) { + auto inputs = ValidSeparateEqualQkvInputs(0); + inputs.query = Shape({0, 8}); + inputs.key = inputs.query; + inputs.value = inputs.query; + inputs.token_offset = Shape({1, sequence_length}); + inputs.cumulative_sequence_length = Shape({2}); + return BuildPackedMultiHeadAttentionProblem(inputs); + }; + + auto boundary_problem = make_problem(kLargestRound64SafeSequenceLength); + ASSERT_TRUE(boundary_problem.status.IsOK()) << boundary_problem.status.message; + boundary_problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + EXPECT_TRUE(GetPackedMultiHeadAttentionWorkspaceRecipe(boundary_problem.problem).status.IsOK()); + + auto beyond_problem = make_problem(kLargestRound64SafeSequenceLength + 1); + ASSERT_TRUE(beyond_problem.status.IsOK()) << beyond_problem.status.message; + beyond_problem.problem.backend = PackedAttentionBackend::MemoryEfficient; + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(beyond_problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, WorkspaceViewsAreContainedInAttentionAllocation) { + auto problem = BuildPackedMultiHeadAttentionProblem(ValidPackedQkvInputs()); + ASSERT_TRUE(problem.status.IsOK()) << problem.status.message; + problem.problem.backend = PackedAttentionBackend::Unfused; + auto workspace = GetPackedMultiHeadAttentionWorkspaceRecipe(problem.problem); + ASSERT_TRUE(workspace.status.IsOK()) << workspace.status.message; + EXPECT_TRUE(ValidatePackedAttentionWorkspaceRecipe(workspace.recipe).IsOK()); + + auto out_of_bounds = workspace.recipe; + out_of_bounds.v_offset_bytes = out_of_bounds.attention_workspace_bytes; + out_of_bounds.v_bytes = 1; + EXPECT_EQ(ValidatePackedAttentionWorkspaceRecipe(out_of_bounds).error, + PackedAttentionWorkspaceError::InvalidArgument); + + out_of_bounds.v_offset_bytes = std::numeric_limits::max(); + EXPECT_EQ(ValidatePackedAttentionWorkspaceRecipe(out_of_bounds).error, + PackedAttentionWorkspaceError::Overflow); + + auto invalid_layout = workspace.recipe; + invalid_layout.qkv_layout = static_cast(999); + EXPECT_EQ(ValidatePackedAttentionWorkspaceRecipe(invalid_layout).error, + PackedAttentionWorkspaceError::InvalidArgument); + + auto missing_second_scratch = workspace.recipe; + missing_second_scratch.has_second_scratch = false; + EXPECT_EQ(ValidatePackedAttentionWorkspaceRecipe(missing_second_scratch).error, + PackedAttentionWorkspaceError::InvalidArgument); + + problem.problem.backend = static_cast(999); + EXPECT_EQ(GetPackedMultiHeadAttentionWorkspaceRecipe(problem.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); +} + +TEST(PackedAttentionWorkspaceTest, PaRejectsFlashAndTrtRequiresExistingRunner) { + auto problem_result = BuildPackedAttentionProblem(ValidPackedAttentionInputs()); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + + problem_result.problem.backend = PackedAttentionBackend::Flash; + EXPECT_EQ(GetPackedAttentionWorkspaceRecipe(problem_result.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + problem_result.problem.backend = PackedAttentionBackend::Trt; + problem_result.problem.trt_runner_available = false; + EXPECT_EQ(GetPackedAttentionWorkspaceRecipe(problem_result.problem).status.error, + PackedAttentionWorkspaceError::InvalidArgument); + + // A missing TRT prerequisite is represented by the actual selected fallback route. + problem_result.problem.backend = PackedAttentionBackend::Unfused; + auto fallback = GetPackedAttentionWorkspaceRecipe(problem_result.problem); + ASSERT_TRUE(fallback.status.IsOK()) << fallback.status.message; + EXPECT_EQ(fallback.recipe.attention_workspace_bytes, 832U); +} + +TEST(PackedAttentionWorkspaceTest, ExactZeroWorkspaceIsValid) { + PackedAttentionInputShapes inputs; + inputs.input = Shape({0, 3}); + inputs.weights = Shape({3, 0}); + inputs.bias = Shape({0}); + inputs.token_offset = Shape({0, 0}); + inputs.cumulative_sequence_length = Shape({1}); + inputs.element_size = 2; + inputs.num_heads = 2; + inputs.qkv_hidden_sizes_count = 3; + inputs.qkv_hidden_sizes = {0, 0, 0}; + + auto problem_result = BuildPackedAttentionProblem(inputs); + ASSERT_TRUE(problem_result.status.IsOK()) << problem_result.status.message; + auto workspace_result = GetPackedAttentionWorkspaceRecipe(problem_result.problem); + ASSERT_TRUE(workspace_result.status.IsOK()) << workspace_result.status.message; + EXPECT_EQ(workspace_result.recipe.projection_bytes, 0U); + EXPECT_EQ(workspace_result.recipe.attention_workspace_bytes, 0U); + EXPECT_TRUE(workspace_result.recipe.has_second_scratch); + EXPECT_EQ(workspace_result.recipe.second_scratch_offset_bytes, 0U); + EXPECT_TRUE(ValidatePackedAttentionWorkspaceRecipe(workspace_result.recipe).IsOK()); +} + +} // namespace test +} // namespace onnxruntime From 33452bf477b400df44c698be019efddc5738af1d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:22:57 +0000 Subject: [PATCH 2/3] Handle empty packed attention outputs Return before CUDA dispatch for empty packed attention outputs and use nonzero padded data in route-observed tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78 --- .../contrib_ops/cuda/bert/packed_attention.cc | 4 ++ .../cuda/bert/packed_multihead_attention.cc | 4 ++ .../contrib_ops/packed_attention_op_test.cc | 58 +++++++++++++++++-- .../packed_multihead_attention_op_test.cc | 25 ++++++++ 4 files changed, 87 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc index 1508d646aa3f0..88bfa1674595e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/packed_attention.cc @@ -190,6 +190,10 @@ Status PackedAttention::ComputeInternal(OpKernelContext* context) const { TensorShapeVector output_shape{parameters.token_count, parameters.v_hidden_size}; Tensor* output = context->Output(0, output_shape); + if (output->Shape().Size() == 0) { + return Status::OK(); + } + auto& device_prop = this->GetDeviceProp(); MHARunner* fused_runner = this->GetFusedRunner(device_prop, attention_bias != nullptr, parameters); diff --git a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc index 5b66b1c35a7c1..8fb8476ffb32c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention.cc @@ -129,6 +129,10 @@ Status PackedMultiHeadAttention::ComputeInternal(OpKernelContext* context) co TensorShapeVector output_shape{parameters.token_count, parameters.v_hidden_size}; Tensor* output = context->Output(0, output_shape); + if (output->Shape().Size() == 0) { + return Status::OK(); + } + auto& device_prop = this->GetDeviceProp(); bool use_flash_attention = false; diff --git a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc index 74122001bd796..1ae779c264fed 100644 --- a/onnxruntime/test/contrib_ops/packed_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_attention_op_test.cc @@ -177,12 +177,37 @@ static void RunPackedAttentionRouteTest(PackedAttentionRoute route) { constexpr int kHeadSize = kHiddenSize / kNumHeads; static_assert(kTokenCount < kBatchSize * kSequenceLength); - const std::vector input_data(kTokenCount * kHiddenSize, 0.0f); - const std::vector weight_data(kHiddenSize * 3 * kHiddenSize, 0.0f); - const std::vector bias_data(3 * kHiddenSize, 0.0f); + // Each token has input [1, x, 0, ...]. The block projection produces + // Q=x+0.25, K=1, and V=2x-0.75 in every head coordinate. Since K is + // constant within each sequence, attention is uniform: the one-token + // sequence returns 1.25 and the two-token sequence averages 3.25 and 7.25. + const std::vector token_values{1.0f, 2.0f, 4.0f}; + std::vector input_data(kTokenCount * kHiddenSize, 0.0f); + std::vector weight_data(kHiddenSize * 3 * kHiddenSize, 0.0f); + std::vector bias_data(3 * kHiddenSize, 0.0f); + std::vector output_data(kTokenCount * kHiddenSize); + const std::vector expected_token_values{1.25f, 5.25f, 5.25f}; + + for (int token = 0; token < kTokenCount; ++token) { + input_data[token * kHiddenSize] = 1.0f; + input_data[token * kHiddenSize + 1] = token_values[token]; + for (int hidden = 0; hidden < kHiddenSize; ++hidden) { + output_data[token * kHiddenSize + hidden] = expected_token_values[token]; + } + } + + constexpr int kProjectionSize = 3 * kHiddenSize; + for (int hidden = 0; hidden < kHiddenSize; ++hidden) { + weight_data[kProjectionSize + hidden] = 1.0f; + weight_data[kHiddenSize + hidden] = 0.5f; + weight_data[kProjectionSize + 2 * kHiddenSize + hidden] = 2.0f; + bias_data[hidden] = 0.25f; + bias_data[kHiddenSize + hidden] = 0.5f; + bias_data[2 * kHiddenSize + hidden] = -0.75f; + } + const std::vector token_offset{0, 2, 3, 1}; const std::vector cumulative_sequence_length{0, 1, 3}; - const std::vector output_data(kTokenCount * kHiddenSize, 0.0f); if (route == PackedAttentionRoute::Trt) { if (!IsTrtFusedAttentionRouteObservable(kHeadSize, kSequenceLength)) { @@ -295,6 +320,31 @@ TEST(PackedAttentionTest, PackedRouteObservedUnfusedWithPadding) { RunPackedAttentionRouteTest(PackedAttentionRoute::Unfused); } +TEST(PackedAttentionTest, EmptyTokensAndSequence_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "PackedAttention empty-output test requires a CUDA device."; + } + + constexpr int kBatchSize = 2; + constexpr int kHiddenSize = 32; + constexpr int kNumHeads = 2; + + OpTester tester("PackedAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", kNumHeads); + tester.AddInput("input", {0, kHiddenSize}, {}); + tester.AddInput("weight", {kHiddenSize, 3 * kHiddenSize}, + std::vector(kHiddenSize * 3 * kHiddenSize, 1.0f)); + tester.AddInput("bias", {3 * kHiddenSize}, + std::vector(3 * kHiddenSize, 0.5f)); + tester.AddInput("token_offset", {kBatchSize, 0}, {}); + tester.AddInput("cumulative_sequence_length", {kBatchSize + 1}, {0, 0, 0}); + tester.AddOutput("output", {0, kHiddenSize}, {}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + TEST(PackedAttentionTest, NoPack) { int batch_size = 1; int sequence_length = 2; diff --git a/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc index 5168f007c6e57..5247170f56728 100644 --- a/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/packed_multihead_attention_op_test.cc @@ -381,6 +381,31 @@ static void RunPackedMultiHeadAttentionTest( } } +TEST(PackedMultiHeadAttentionTest, EmptyTokensAndSequence_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "PackedMultiHeadAttention empty-output test requires a CUDA device."; + } + + constexpr int kBatchSize = 2; + constexpr int kNumHeads = 2; + constexpr int kHeadSize = 16; + constexpr int kHiddenSize = kNumHeads * kHeadSize; + + OpTester tester("PackedMultiHeadAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", kNumHeads); + tester.AddInput("query", {0, kNumHeads, 3, kHeadSize}, {}); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddInput("token_offset", {kBatchSize, 0}, {}); + tester.AddInput("cumulative_sequence_length", {kBatchSize + 1}, {0, 0, 0}); + tester.AddOutput("output", {0, kHiddenSize}, {}); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + TEST(PackedMultiHeadAttentionTest, PackedQKV_NoPadding_NoBias_trt) { AttentionTestData data; GetSelfAttentionData_Batch2_HeadSize32_NoBias_NoMask_PackedQKV(data); From 3b4666905ab2b591cf34a9ddbc9aa2074caf7316 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:13:09 +0000 Subject: [PATCH 3/3] Document attention workspace estimation roadmap Describe the operator-specific rollout, runtime versus AOT sizing contract, and packed attention PR1 scope alongside the CUDA workspace inventory. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78 --- .../attention_workspace_estimation.md | 172 ++++++++++++++++++ .../cuda_kernel_workspace_inventory.md | 60 ++++-- 2 files changed, 215 insertions(+), 17 deletions(-) create mode 100644 docs/annotated_partitioning/attention_workspace_estimation.md diff --git a/docs/annotated_partitioning/attention_workspace_estimation.md b/docs/annotated_partitioning/attention_workspace_estimation.md new file mode 100644 index 0000000000000..9a0118e099744 --- /dev/null +++ b/docs/annotated_partitioning/attention_workspace_estimation.md @@ -0,0 +1,172 @@ +# CUDA Attention Workspace Estimation Roadmap + +## Goal and scope + +This work provides operator-specific workspace estimation for the CUDA Attention family. It defines how an +Attention kernel derives a workspace recipe from a plain problem description and a selected backend. + +The generic memory-estimation framework is a separate work track and is out of scope here. That track includes L1/L2 +callback plumbing, kernel registration, shape and max-shape infrastructure, optional-input contracts at the framework +boundary, partition budgeting, memory planning, and planned-root or preallocation integration. + +The Attention work owns: + +- graph-free problem descriptions and checked workspace recipes; +- backend- and layout-specific workspace formulas; +- runtime allocation and view parity; +- route and boundary tests; and +- thin Attention-specific adapters after the generic framework API stabilizes. + +## Runtime dispatch and AOT estimation + +CUDA EP assignment happens during session initialization. Backend dispatch inside an assigned CUDA Attention kernel +can happen later, using the concrete inputs available to `Compute()`. These are distinct decisions. + +Runtime sizing can be exact because the kernel first selects a backend and then requests that backend's recipe: + +```text +concrete inputs -> runtime dispatch -> selected backend -> exact workspace recipe +``` + +Ahead-of-time (AOT) estimation cannot assume that the same selection is known. Backend feasibility can depend on +optional inputs, build and runtime options, cache state, runner availability, device properties, and concrete dynamic +sequence lengths. When the exact backend cannot be proven, AOT estimation must enumerate the feasible backend recipes +and take a safe upper bound: + +```text +shapes and bounds -> feasible backends -> recipe per backend -> maximum workspace +``` + +An estimator must not copy the runtime dispatch cascade or assume that dispatch is monotonic with shape. For example, +a larger shape may use a fused backend with small workspace while a nearby smaller shape falls back to an unfused +backend with an `S^2` attention buffer. + +An AOT result should therefore be classified as: + +- **Exact:** the backend and all governing dimensions are proven. +- **Safe bound:** the maximum workspace across all feasible backend recipes. +- **Unavailable:** a required shape, optional-input, capability, or recipe contract is not available. + +## Recipe architecture + +The reusable implementation follows this boundary: + +```text +operator shapes and attributes + | + v +plain Attention problem + | + v +checked backend/layout recipe + | + v +runtime allocation and workspace views +``` + +The plain problem and recipe must not depend on `Node`, `NodeArg`, `GraphViewer`, `TensorShape`, or CUDA runtime types. +A future framework adapter may translate framework inputs into the plain problem, but it must not duplicate sizing or +validation arithmetic. + +Recipes must use checked arithmetic, enforce relevant CUDA ABI and grid limits, and prove that every derived view is +contained in the allocated workspace. + +## PR1: packed Attention recipes + +[microsoft/onnxruntime#32283](https://github.com/microsoft/onnxruntime/pull/32283) is the first implementation in the +roadmap tracked by [microsoft/onnxruntime#29775](https://github.com/microsoft/onnxruntime/issues/29775). Its scope is +`PackedAttention` (PA) and `PackedMultiHeadAttention` (PMHA), plus a shared correction that widens CUTLASS MEA +attention-bias stride arithmetic for all MEA consumers. + +PR1 establishes a single sizing and layout source of truth while preserving legacy workspace byte totals, allocation +counts, and allocation lifetimes. It does not change Attention backend selection. The shared stride correction can +change MEA's internal aligned-versus-unaligned kernel variant only in cases where the previous int32 calculation +overflowed before assignment to an int64 stride. + +### `T` and `B * S` + +`T` is the packed real-token count. `B * S` is padded capacity. Existing packed Attention paths use both: + +- fused backend views can be governed by `T`; +- the unfused Q/K/V layout is governed by `B * S`; and +- PR1 retains the legacy `B * S` attention allocation total even when a fused inner view uses `T`. + +Shrinking the total allocation from `B * S` to `T` is a separate optimization. + +### Workspace components and layouts + +PA has a projection GEMM allocation and an Attention allocation. The recipe reports them separately; their sum must +not replace either allocation. PMHA receives Q/K/V inputs and has zero projection workspace. + +| Backend | Q/K/V representation | Q/K/V view dimension | Backend scratch retained by PR1 | +| --- | --- | --- | --- | +| Flash | Planar materialized views or direct input views | `T` | Softmax LSE: `sizeof(float) * B * S * N` | +| TensorRT fused (`FusedRunner`) | Interleaved `[T, N, 3, H]` | `T` | None | +| Memory-efficient Attention | Planar materialized views or direct input views | `T` | Optional FP32 accumulator: `sizeof(float) * B * S * N * H_v` | +| Unfused (`Default`) | Planar `[B, N, S, H]` views | `B * S` | Two individually aligned `element_size * B * N * S * S` regions | + +The MEA accumulator is needed when `H_v > 128` and the input element size is smaller than FP32. PA does not dispatch +to Flash and always materializes Q/K/V after its projection GEMM. PMHA can skip materialization for packed +`[T, N, 3, H]` input on TensorRT when bias is absent, or for separate Q/K/V input on Flash or MEA when bias is absent. + +### Validation and tests + +PR1 includes: + +- checked size, offset, alignment, and derived-stride arithmetic; +- CUDA int32 and int64 ABI validation scoped to the selected backend and materialization producer; +- explicit planar, interleaved, and direct-view contracts; +- recipe containment validation; +- independent hand-calculated byte and layout parity tests; +- runtime route tests for Flash, TensorRT fused, memory-efficient, and unfused paths as applicable; and +- empty-output handling before GEMM or CUDA kernel dispatch. + +PR1 validates token-offset and cumulative-sequence tensor shapes and host-visible geometry. It does not inspect or +synchronize their device contents. Device-value validation requires a separate CUDA graph- and capture-safe contract. + +### PR1 non-goals + +PR1 does not: + +- connect to the generic L1/L2 framework; +- change Attention backend selection or eligibility; +- change allocation count or lifetime; +- shrink the legacy `B * S` total to `T`; +- introduce planned-root allocation or preallocation; or +- validate device-side token-offset or cumulative-sequence values. + +## Attention family rollout + +PR1 precedes the sequence below and establishes the checked-recipe architecture that later families can reuse. + +The planned operator-specific sequence is: + +1. Shared backend primitives and `MultiHeadAttention`. +2. `GroupQueryAttention`. +3. `PagedAttention`. +4. High-value decoder-specific variants. +5. Separate memory models for Linear, Sparse, Longformer, quantized, and ONNX Attention operators. +6. Thin Attention-specific adapters to the generic framework after its API stabilizes. + +MHA and GQA are high-value coverage targets and have high estimation-drift risk. Their runtime behavior can include +dynamic internal backend dispatch, cache lifecycle and aliasing, optional inputs, non-monotonic fallback paths, and +unfused workspace governed by `S_q * S_kv_total`. GQA additionally has different Q and KV head counts. MHA can have +different query and total-KV sequence lengths or different Q/K and V head sizes. Shared backend kernels do not +eliminate operator-specific preparation, cache, transpose, grouping, or output workspace. + +Paged, Linear, Sparse, and Longformer Attention require distinct memory models. They must not be forced into a dense +Attention formula solely because they share some backend infrastructure. + +## Acceptance standard + +Each Attention family must provide: + +- a runtime source of truth for the selected backend's workspace and layout; +- checked arithmetic and ABI, grid, alignment, and containment validation; +- byte-total and view-layout parity with runtime allocation; +- tests for feasible backend routes and fallback boundaries; +- explicit exact, safe-bound, or unavailable estimation semantics; +- no copied runtime dispatch cascade; and +- no dependency from the reusable recipe on graph or CUDA runtime types. + +This document defines operator-estimation work tracks and invariants. It does not prescribe a public framework API. diff --git a/docs/annotated_partitioning/cuda_kernel_workspace_inventory.md b/docs/annotated_partitioning/cuda_kernel_workspace_inventory.md index bb6a73568124b..d33c419748890 100644 --- a/docs/annotated_partitioning/cuda_kernel_workspace_inventory.md +++ b/docs/annotated_partitioning/cuda_kernel_workspace_inventory.md @@ -8,6 +8,7 @@ This document catalogs all CUDA kernels in ONNX Runtime that allocate temporary/ |--------|---------| | ✅ | Fully determinable from shapes + attributes + device properties | | ✅* | Determinable via cuDNN/cuBLAS API call (needs handle, available on EP) | +| 🔀 | Runtime-exact; AOT exact only when backend selection is provable | | ⚠️ | Requires profiling/tactic selection (deterministic but costly at planning time) | --- @@ -299,18 +300,37 @@ This document catalogs all CUDA kernels in ONNX Runtime that allocate temporary/ ### 18. Attention / MultiHeadAttention (Contrib) -**File:** `bert/attention.cc`, `bert/multihead_attention.cc` +**File:** `bert/attention.cc`, `bert/multihead_attention.cc`, `bert/packed_attention.cc`, +`bert/packed_multihead_attention.cc` -**Buffers:** Uses `GetAttentionWorkspaceSize()` helper function. +**Roadmap:** [CUDA Attention workspace estimation](attention_workspace_estimation.md) -**Size formula:** Depends on attention algorithm (Flash, MemoryEfficient, FusedRunner, Default): -- Flash: `qkv_size` (Q+K+V projection) -- MemoryEfficient: `qkv_size + output_accum (float)` -- Default (unfused): `qkv_size + 2 * attention_scratch_size` +**Buffers:** Dense Attention uses the `GetAttentionWorkspaceSize()` helper. MultiHeadAttention can additionally +allocate lean-attention synchronization, Flash split/LSE/output-accumulator, and sequence-length buffers. +PackedAttention has separate projection and Attention allocations; PackedMultiHeadAttention has no projection +allocation. See the roadmap for the packed recipe and layout contract. -**What's needed:** B, S_q, S_kv, num_heads, head_size, dtype, which attention algorithm is selected. +**Size model:** Depends on operator inputs and the selected Attention algorithm: -**Static determinability:** ✅ — Algorithm selection depends on shapes + SM version (available from device_prop). +- Dense Attention/MHA: the main helper covers backend-dependent Q/K/V materialization and Attention scratch. + MHA computes lean synchronization, Flash split/LSE/output-accumulator, and sequence-length allocations separately. +- PackedAttention/PMHA Flash: optional planar Q/K/V materialization plus + `sizeof(float) * B * S * num_heads` for softmax LSE. +- PackedAttention/PMHA TensorRT fused: optional interleaved `[T, N, 3, H]` QKV materialization and no backend scratch. +- PackedAttention/PMHA MemoryEfficient: optional planar Q/K/V materialization plus an optional + `sizeof(float) * B * S * num_heads * v_head_size` accumulator. +- PackedAttention/PMHA unfused: `B * S` planar Q/K/V capacity plus two individually aligned + `element_size * B * num_heads * S * S` scratch regions. + +PA always materializes Q/K/V after its projection allocation. PMHA can use direct input views for eligible fused +routes without bias. See the roadmap for the precise layout and eligibility boundaries. + +**What's needed:** B, S_q, S_kv, num_heads, head_size, dtype, selected Attention algorithm, optional-input presence, +cache state, and device properties including SM version and `multiProcessorCount`. + +**Static determinability:** 🔀 — Runtime-exact when sizing receives the backend selected by the CUDA kernel. For AOT, +use the maximum workspace across feasible backend recipes when exact selection is not provable, or report estimation +as unavailable when a required contract is missing. See the linked roadmap for the exact/safe-bound/unavailable model. --- @@ -403,8 +423,9 @@ GEMM); see that entry. Not re-verified independently of MatMulNBits's call sites | Category | # Kernels | Estimation feasibility | Notes | |----------|-----------|----------------------|-------| -| **Shapes only** | 12 | ✅ Exact, trivial | BatchNorm, InstanceNorm, Dropout, TopK, MatMulInteger, IntegerGemm, Compress, GatherND, NonZero, Upsample, Inverse, Generation | -| **Shapes + device properties** | 3 | ✅ Exact | Attention (SM count), DeformConv (totalGlobalMem), Contrib Attention (SM version) | +| **Shapes only** | 13 | ✅ Exact, trivial | BatchNorm, InstanceNorm, Dropout, TopK, MatMulInteger, IntegerGemm, Compress, GatherND, NonZero, Upsample, NonMaxSuppression, Inverse, Generation | +| **Shapes + device properties** | 2 | ✅ Exact | Attention (SM count), DeformConv (totalGlobalMem) | +| **Shapes + backend feasibility** | 1 | 🔀 Runtime-exact, AOT conditional | Contrib Attention | | **Shapes + cuDNN/cuBLAS handle** | 4 | ✅* Exact via API query | Conv, ConvTranspose, Reduction, RNN | | **Shapes + closed-form GEMM formula** | 2 | ✅ Exact (confirmed PR #29811) | MatMulNBits, fpA_intB_GEMM | | **Shapes + tactic profiling** | 1 | ⚠️† Upper bound only | MOE | @@ -414,9 +435,12 @@ do not assume it shares their exact-formula property. ### Key takeaways -1. **~84% of kernels** (21/25) can produce **exact** workspace estimates at `GetCapability()` time using only shapes + attributes + device properties (+ cuDNN handle for API queries, or a closed-form GEMM formula for the fpA_intB pair). +1. **~91% of kernels** (21/23) can produce **exact** workspace estimates at `GetCapability()` time using only shapes + attributes + device properties (+ cuDNN handle for API queries, or a closed-form GEMM formula for the fpA_intB pair). + +2. **~4% of kernels** (1/23, Contrib Attention) are runtime-exact but AOT-conditional. AOT estimation must prove the + backend or use the maximum across feasible backend recipes. -2. **~4% of kernels** (1/25, MOE) require tactic profiling (CUTLASS/CUB autotuning). For this kernel, options are: +3. **~4% of kernels** (1/23, MOE) require tactic profiling (CUTLASS/CUB autotuning). For this kernel, options are: - Use worst-case workspace across all tactics (safe upper bound) - Run tactic selection eagerly at estimation time (expensive but exact) - Accept a safety multiplier for this one kernel @@ -424,12 +448,14 @@ do not assume it shares their exact-formula property. MatMulNBits and fpA_intB_GEMM were previously grouped here too; see entries #20/#21 above for why PR #29811 moved them to the closed-form-exact row instead. -3. **The cuDNN handle requirement** affects only 4 kernel types (Conv, ConvTranspose, Reduction, RNN). All are standard cuDNN API queries that are fast and deterministic given the handle + tensor descriptors. +4. **The cuDNN handle requirement** affects only 4 kernel types (Conv, ConvTranspose, Reduction, RNN). All are standard cuDNN API queries that are fast and deterministic given the handle + tensor descriptors. -4. **No kernel requires actual GPU execution** to determine workspace size — even tactic-based kernels select tactics via CPU-side profiling/heuristics, not by running GPU code. +5. **Most kernels do not require actual GPU execution** to determine workspace size. MOE exact tactic selection may + require GPU-side profiling; use a safe upper bound when that profiling is not appropriate during planning. -5. **Largest workspace consumers** in practice: - - **Attention** (Flash): dominates in LLM workloads. Exact estimation possible. +6. **Largest workspace consumers** in practice: + - **Attention**: dominates in LLM workloads. Runtime sizing can be exact; AOT estimation is conditional on backend + feasibility and may require a safe upper bound. - **Conv** (cuDNN): dominates in vision workloads. Exact via `build_plans()`. - **MOE**: significant in MoE models. Upper bound via worst-case tactic. @@ -437,7 +463,7 @@ do not assume it shares their exact-formula property. | Access needed | How accessed | Kernels that require it | |---------------|-------------|------------------------| -| `Node_GetInputShape()` | OrtEpApi (generic) | All 25 kernels | +| `Node_GetInputShape()` | OrtEpApi (generic) | All 23 kernels | | `Node_GetAttributeInt/Ints()` | OrtEpApi (generic) | Conv, Attention, RNN, MOE | | `device_prop.multiProcessorCount` | Cast `OrtEp*` to concrete EP type | Attention, DeformConv, MatMulNBits, fpA_intB | | `device_prop.totalGlobalMem` | Cast `OrtEp*` to concrete EP type | DeformConv |