From d66c7a839934035fc873af76e306832e0c1df552 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 16 Jul 2026 23:07:49 -0700 Subject: [PATCH 01/20] [QNN EP] Add support for Attention Op --- onnxruntime/core/providers/qnn/builder/op_builder_factory.cc | 1 + onnxruntime/core/providers/qnn/builder/op_builder_factory.h | 1 + 2 files changed, 2 insertions(+) diff --git a/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc b/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc index 3781bff20a6..7b3ae85b6a7 100644 --- a/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc +++ b/onnxruntime/core/providers/qnn/builder/op_builder_factory.cc @@ -13,6 +13,7 @@ namespace qnn { OpBuilderRegistrations::OpBuilderRegistrations() { CreateArgMaxMinOpBuilder("ArgMax", *this); CreateArgMaxMinOpBuilder("ArgMin", *this); + CreateAttentionOpBuilder("Attention", *this); CreateBatchNormalizationOpBuilder("BatchNormalization", *this); CreateCastOpBuilder("Cast", *this); CreateClipOpBuilder("Clip", *this); diff --git a/onnxruntime/core/providers/qnn/builder/op_builder_factory.h b/onnxruntime/core/providers/qnn/builder/op_builder_factory.h index 31fc2605678..3a229bea7f9 100644 --- a/onnxruntime/core/providers/qnn/builder/op_builder_factory.h +++ b/onnxruntime/core/providers/qnn/builder/op_builder_factory.h @@ -58,6 +58,7 @@ class OpBuilderRegistrations { const IOpBuilder* GetOpBuilder(const std::string& onnx_op_type); void CreateArgMaxMinOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateBatchNormalizationOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateCastOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateClipOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); From a0a2ce3bfbb65b48a81ee7c0ba722dfb67f5b71f Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Fri, 17 Jul 2026 12:18:20 -0700 Subject: [PATCH 02/20] [QNN EP] Add support for Attention Op --- .../builder/opbuilder/attention_op_builder.cc | 1139 +++++++++++++++++ .../test/providers/qnn/attention_test.cc | 554 ++++++++ 2 files changed, 1693 insertions(+) create mode 100644 onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc create mode 100644 onnxruntime/test/providers/qnn/attention_test.cc diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc new file mode 100644 index 00000000000..9255b25ceb3 --- /dev/null +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -0,0 +1,1139 @@ +// Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include + +#include "QnnOpDef.h" +#include "core/providers/qnn/builder/op_builder_factory.h" +#include "core/providers/qnn/builder/opbuilder/base_op_builder.h" +#include "core/providers/qnn/builder/qnn_model_wrapper.h" +#include "core/providers/qnn/builder/qnn_utils.h" + +namespace onnxruntime { +namespace qnn { + +// Decomposition builder for ai.onnx::Attention (opset 23/24). +// +// 4D inputs [B, n, S, hs] (BNSH layout, no reshape needed) +// 3D inputs [B, S, n*hs] (BSH layout, reshape + transpose to BNSH) +// GQA/MQA, KV cache (past/present), softcap, qk_matmul_output +// +// The SDK version guard matches GroupQueryAttentionOpBuilder so that this class +// is compiled away when building against an SDK without the required op-set +// version metadata. +#if !(QNN_OPSET_VERSION_MAJOR < 2 || (QNN_OPSET_VERSION_MAJOR == 2 && QNN_OPSET_VERSION_MINOR <= 11)) + +class AttentionOpBuilder : public BaseOpBuilder { + public: + AttentionOpBuilder() : BaseOpBuilder("AttentionOpBuilder") {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(AttentionOpBuilder); + + protected: + Ort::Status IsOpSupported(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger) const override ORT_MUST_USE_RESULT; + + Ort::Status ProcessInputs(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger, + std::vector& input_names, + bool do_op_validation) const override ORT_MUST_USE_RESULT; + + Ort::Status ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + std::vector&& input_names, + const Ort::Logger& logger, + bool do_op_validation) const override ORT_MUST_USE_RESULT; +}; + +// --------------------------------------------------------------------------- +// IsOpSupported +// --------------------------------------------------------------------------- +Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger) const { + ORT_UNUSED_PARAMETER(logger); + + const auto& inputs = node_unit.Inputs(); + const auto& outputs = node_unit.Outputs(); + const size_t num_inputs = inputs.size(); + const size_t num_outputs = outputs.size(); + + OrtNodeAttrHelper node_helper(node_unit); + + // ---- Primary output must exist ---- + RETURN_IF_NOT(num_outputs > 0 && outputs[0].Exists(), + "Required output Y (output[0]) not provided"); + + // ---- KV cache: past_key and present_key must both be present or both absent ---- + const bool has_past_key = (num_inputs > 4 && inputs[4].Exists()); + const bool has_past_value = (num_inputs > 5 && inputs[5].Exists()); + const bool has_present_key = (num_outputs > 1 && outputs[1].Exists()); + const bool has_present_value = (num_outputs > 2 && outputs[2].Exists()); + + RETURN_IF(has_past_key != has_past_value, + "Attention: past_key and past_value must both be present or both absent"); + RETURN_IF(has_present_key != has_present_value, + "Attention: present_key and present_value must both be present or both absent"); + RETURN_IF(has_past_key && !has_present_key, + "Attention: past_key present but present_key output not provided"); + + // ---- nonpad_kv_seqlen (input[6], opset 24) ---- + const bool has_nonpad_kv_seqlen = (num_inputs > 6 && inputs[6].Exists()); + RETURN_IF(has_nonpad_kv_seqlen && has_past_key, + "Attention: nonpad_kv_seqlen and past_key are mutually exclusive per ONNX spec"); + + // ---- qk_matmul_output_mode ---- + const int64_t qk_mode = node_helper.Get("qk_matmul_output_mode", static_cast(0)); + RETURN_IF(qk_mode < 0 || qk_mode > 3, + "Attention: qk_matmul_output_mode must be in [0,3]"); + // If mode != 0 but output[3] is not provided, we just ignore (not an error). + // If output[3] is provided but mode == 0, it's mode-0 (post QK matmul). + + // ---- scale must be positive if provided ---- + if (node_helper.HasAttr("scale")) { + const float scale_check = node_helper.Get("scale", 1.0f); + RETURN_IF(scale_check <= 0.0f, "scale attribute must be positive"); + } + + // ---- Validate Q shape and determine input layout ---- + TensorInfo q_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[0], q_info)); + const size_t q_rank = q_info.shape.size(); + RETURN_IF(q_rank != 3 && q_rank != 4, + "Attention: Q input must be rank 3 ([B,S_q,n_q*hs]) or rank 4 ([B,n_q,S_q,hs])"); + + // ---- Reject dynamic (0-dim) shapes on Q, K, V ---- + for (uint32_t d : q_info.shape) { + RETURN_IF(d == 0, + "Attention: Q input contains a dynamic (0) dimension; only static shapes are supported"); + } + + TensorInfo k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[1], k_info)); + for (uint32_t d : k_info.shape) { + RETURN_IF(d == 0, + "Attention: K input contains a dynamic (0) dimension; only static shapes are supported"); + } + + TensorInfo v_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[2], v_info)); + for (uint32_t d : v_info.shape) { + RETURN_IF(d == 0, + "Attention: V input contains a dynamic (0) dimension; only static shapes are supported"); + } + + // ---- For 3D inputs, q_num_heads and kv_num_heads attrs are required ---- + if (q_rank == 3) { + RETURN_IF_NOT(node_helper.HasAttr("q_num_heads"), + "Attention: q_num_heads attribute required for 3D (BSH) inputs"); + RETURN_IF_NOT(node_helper.HasAttr("kv_num_heads"), + "Attention: kv_num_heads attribute required for 3D (BSH) inputs"); + } + + // ---- GQA divisibility check ---- + { + uint32_t q_nh = 0; + uint32_t kv_nh = 0; + if (q_rank == 4) { + q_nh = q_info.shape[1]; + kv_nh = k_info.shape[1]; + } else { + const auto opt_q = node_helper.GetInt64("q_num_heads"); + const auto opt_kv = node_helper.GetInt64("kv_num_heads"); + if (opt_q.has_value()) q_nh = static_cast(opt_q.value()); + if (opt_kv.has_value()) kv_nh = static_cast(opt_kv.value()); + } + RETURN_IF(q_nh != 0 && kv_nh != 0 && q_nh != kv_nh && q_nh % kv_nh != 0, + "Attention: GQA requires q_num_heads to be divisible by kv_num_heads"); + } + + // ---- KV cache: require static S_past ---- + if (has_past_key) { + TensorInfo past_k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[4], past_k_info)); + RETURN_IF(past_k_info.shape.size() != 4, + "Attention: past_key must be rank 4 ([B,n,S_past,hs])"); + for (uint32_t d : past_k_info.shape) { + RETURN_IF(d == 0, + "Attention: past_key contains a dynamic (0) dimension; only static shapes are supported"); + } + } + + // ---- Full validation: build decomposed nodes with do_op_validation=true ---- + std::vector input_names; + RETURN_IF_ERROR(ProcessInputs(qnn_model_wrapper, node_unit, logger, input_names, true)); + RETURN_IF_ERROR( + ProcessAttributesAndOutputs(qnn_model_wrapper, node_unit, std::move(input_names), logger, true)); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// ProcessInputs — register Q, K, V, attn_mask, past_key, past_value, +// nonpad_kv_seqlen +// --------------------------------------------------------------------------- +Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger, + std::vector& input_names, + bool /*do_op_validation*/) const { + const auto& onnx_inputs = node_unit.Inputs(); + + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[0], logger, input_names)); // Q + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[1], logger, input_names)); // K + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[2], logger, input_names)); // V + + // input[3] = attn_mask (optional) + if (onnx_inputs.size() > 3 && onnx_inputs[3].Exists()) { + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[3], logger, input_names)); + } + // input[4] = past_key (optional, KV cache) + if (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()) { + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); + } + // input[5] = past_value (optional, KV cache) + if (onnx_inputs.size() > 5 && onnx_inputs[5].Exists()) { + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); + } + // input[6] = nonpad_kv_seqlen (optional, opset 24) + if (onnx_inputs.size() > 6 && onnx_inputs[6].Exists()) { + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[6], logger, input_names)); + } + + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: emit an ElementWiseBinary (MUL or ADD or DIV) node. +// --------------------------------------------------------------------------- +static Ort::Status AddBinaryOpNode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + uint32_t operation, + const std::string& lhs_name, + const std::string& rhs_name, + const std::string& out_name, + const std::vector& out_shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool is_graph_output, + bool do_op_validation) { + const Qnn_TensorType_t tensor_type = + is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + QnnTensorWrapper out_tensor(out_name, tensor_type, dtype, quant_param.Copy(), + std::vector(out_shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), + ("Failed to add output tensor: " + out_name).c_str()); + + const std::string node_name = utils::UniqueNameGenerator().New(node_unit, "_ewb"); + + std::vector param_names; + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, + node_unit.Index(), + node_name, + operation, + QNN_OP_ELEMENT_WISE_BINARY_PARAM_OPERATION, + param_names)); + + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_ELEMENT_WISE_BINARY, + {lhs_name, rhs_name}, + {out_name}, + std::move(param_names), + do_op_validation), + "Failed to create ElementWiseBinary node."); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: emit a MatMul node (with optional transpose_in1). +// --------------------------------------------------------------------------- +static Ort::Status AddMatMulNode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const std::string& lhs_name, + const std::string& rhs_name, + const std::string& out_name, + const std::vector& out_shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool transpose_in1, + bool do_op_validation) { + QnnTensorWrapper out_tensor(out_name, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), + std::vector(out_shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), + ("Failed to add MatMul output tensor: " + out_name).c_str()); + + const std::string node_name = utils::UniqueNameGenerator().New(node_unit, "_matmul"); + + std::vector param_names; + if (transpose_in1) { + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, + node_unit.Index(), + node_name, + true, + QNN_OP_MAT_MUL_PARAM_TRANSPOSE_IN1, + param_names)); + } + + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_MAT_MUL, + {lhs_name, rhs_name}, + {out_name}, + std::move(param_names), + do_op_validation), + "Failed to create MatMul node."); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: emit a Softmax node (axis param). +// --------------------------------------------------------------------------- +static Ort::Status AddSoftmaxNode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const std::string& in_name, + const std::string& out_name, + const std::vector& shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + uint32_t axis, + bool do_op_validation) { + QnnTensorWrapper out_tensor(out_name, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), + std::vector(shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), + ("Failed to add Softmax output tensor: " + out_name).c_str()); + + const std::string node_name = utils::UniqueNameGenerator().New(node_unit, "_softmax"); + + std::vector param_names; + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, + node_unit.Index(), + node_name, + axis, + QNN_OP_SOFTMAX_PARAM_AXIS, + param_names)); + + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_SOFTMAX, + {in_name}, + {out_name}, + std::move(param_names), + do_op_validation), + "Failed to create Softmax node."); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: create a static scalar tensor (fp32 or fp16) for softcap arithmetic. +// --------------------------------------------------------------------------- +static Ort::Status AddScalarTensor(QnnModelWrapper& qnn_model_wrapper, + const std::string& name, + float value, + Qnn_DataType_t dtype) { + std::vector bytes; + if (dtype == QNN_DATATYPE_FLOAT_16) { + const Ort::Float16_t fp16(value); + bytes.resize(sizeof(uint16_t)); + const uint16_t raw = fp16.val; + std::memcpy(bytes.data(), &raw, sizeof(uint16_t)); + } else { + bytes.resize(sizeof(float)); + std::memcpy(bytes.data(), &value, sizeof(float)); + } + QnnTensorWrapper t(name, QNN_TENSOR_TYPE_STATIC, dtype, QnnQuantParamsWrapper{}, + {1u}, std::move(bytes)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(t)), + ("Failed to add scalar tensor: " + name).c_str()); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: Softcap node. +// out = softcap * tanh(scores / softcap) +// Steps: Div(scores, sc) → Tanh → Mul(result, sc) +// --------------------------------------------------------------------------- +static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const std::string& in_name, + const std::string& out_name, + const std::vector& shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + float softcap_val, + bool do_op_validation) { + // Static scalar for softcap value. + const std::string sc_name = utils::UniqueNameGenerator().New(node_unit, "_softcap_scalar"); + RETURN_IF_ERROR(AddScalarTensor(qnn_model_wrapper, sc_name, softcap_val, dtype)); + + // Div(scores, softcap) -> x + const std::string div_out = utils::UniqueNameGenerator().New(node_unit, "_softcap_div"); + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_DIVIDE, + in_name, sc_name, div_out, + shape, dtype, quant_param, + /*is_graph_output=*/false, do_op_validation)); + + // Tanh(x) -> t + const std::string tanh_out = utils::UniqueNameGenerator().New(node_unit, "_softcap_tanh"); + { + QnnTensorWrapper tanh_tensor(tanh_out, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), + std::vector(shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(tanh_tensor)), + ("Failed to add softcap Tanh output: " + tanh_out).c_str()); + const std::string tanh_node = utils::UniqueNameGenerator().New(node_unit, "_tanh"); + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(tanh_node, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_TANH, + {div_out}, + {tanh_out}, + {}, + do_op_validation), + "Failed to create softcap Tanh node."); + } + + // Mul(t, softcap) -> out + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + tanh_out, sc_name, out_name, + shape, dtype, quant_param, + /*is_graph_output=*/false, do_op_validation)); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: GQA head expansion. +// in_shape = [B, n_kv, S, hs] (K or V already in BNSH after 3D→BNSH transform) +// out_shape = [B, n_q, S, hs] +// +// Three-step expansion (correct floor-division semantics): +// 1. Reshape in → [B, n_kv, 1, S, hs] (insert new dim at axis 2) +// 2. Tile [1,1,head_ratio,1,1] → [B, n_kv, head_ratio, S, hs] +// 3. Reshape → [B, n_q, S, hs] +// +// --------------------------------------------------------------------------- +static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const std::string& in_name, + const std::string& out_name, + const std::vector& in_shape, // [B, n_kv, S, hs] + const std::vector& out_shape, // [B, n_q, S, hs] + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + uint32_t head_ratio, + bool /*do_op_validation*/) { + // 4D-only GQA expansion that avoids 5D tensors (HTP finalization fails with 5D). + // + // Goal: produce K_expanded[b, kv*head_ratio+r, s, h] = K[b, kv, s, h] + // = floor-division [K0,K0,...,K1,K1,...] matching ONNX spec. + // + // Steps (all 4D): + // [B, n_kv, S, hs] + // → Reshape [B, 1, n_kv, S*hs] (insert unit dim, merge S+hs) + // → Tile [1, head_ratio, 1, 1] (block-repeat on size-1 dim → safe copies) + // → [B, head_ratio, n_kv, S*hs] + // → Transpose (0,2,1,3) → [B, n_kv, head_ratio, S*hs] + // → Reshape [B, n_q, S, hs] (C-order: kv*head_ratio+r → floor-div ✓) + + const uint32_t B = in_shape[0]; + const uint32_t n_kv = in_shape[1]; + const uint32_t S = in_shape[2]; + const uint32_t hs = in_shape[3]; + const uint32_t Shs = S * hs; // merged dim + + // Step 1: Reshape [B, n_kv, S, hs] → [B, 1, n_kv, S*hs] + const std::string r1_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_r1"); + const std::vector r1_shape = {B, 1u, n_kv, Shs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(in_name, r1_name, + in_shape, r1_shape, + dtype, quant_param, + /*do_op_validation=*/false, + /*is_for_input=*/false)); + + // Step 2: Tile [1, head_ratio, 1, 1] → [B, head_ratio, n_kv, S*hs] + // Block-repeat on size-1 dim 1 is always correct regardless of Tile semantics. + const std::string tiled_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_tile"); + const std::vector tiled_shape = {B, head_ratio, n_kv, Shs}; + { + QnnTensorWrapper tiled_tensor(tiled_name, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), + std::vector(tiled_shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(tiled_tensor)), + ("Failed to add GQA Tile output tensor: " + tiled_name).c_str()); + + const std::string tile_node = utils::UniqueNameGenerator().New(node_unit, "_gqa_tilenode"); + std::vector mult_data = {1u, head_ratio, 1u, 1u}; + QnnParamWrapper mult_param(node_unit.Index(), tile_node, QNN_OP_TILE_PARAM_MULTIPLES, + {4u}, std::move(mult_data)); + std::vector tile_params; + tile_params.push_back(mult_param.GetParamTensorName()); + RETURN_IF_NOT(qnn_model_wrapper.AddParamWrapper(std::move(mult_param)), + "Failed to add GQA Tile multiples param."); + + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(tile_node, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_TILE, + {r1_name}, + {tiled_name}, + std::move(tile_params), + /*do_op_validation=*/false), + "Failed to create GQA Tile node."); + } + + // Step 3: Transpose (0,2,1,3) → [B, n_kv, head_ratio, S*hs] + const std::string tr_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_tr"); + const std::vector tr_shape = {B, n_kv, head_ratio, Shs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + tiled_name, tr_name, + tiled_shape, + {0u, 2u, 1u, 3u}, + tr_shape, + dtype, quant_param, + /*do_op_validation=*/false, + /*is_for_input=*/false)); + + // Step 4: Reshape [B, n_kv, head_ratio, S*hs] → [B, n_q, S, hs] + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(tr_name, out_name, + tr_shape, out_shape, + dtype, quant_param, + /*do_op_validation=*/false, + /*is_for_input=*/false)); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: KV concat (past || current along the sequence axis). +// past_shape = [B, n, S_past, hs] +// cur_shape = [B, n, S_cur, hs] +// out_shape = [B, n, S_past+S_cur, hs] +// --------------------------------------------------------------------------- +static Ort::Status AddKVConcatNode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const std::string& past_name, + const std::string& cur_name, + const std::string& out_name, + const std::vector& out_shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool is_graph_output, + bool do_op_validation) { + const Qnn_TensorType_t tensor_type = + is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + QnnTensorWrapper out_tensor(out_name, tensor_type, dtype, quant_param.Copy(), + std::vector(out_shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), + ("Failed to add KV concat output tensor: " + out_name).c_str()); + + const std::string node_name = utils::UniqueNameGenerator().New(node_unit, "_kv_concat"); + std::vector param_names; + // axis = 2 (the sequence dimension in BNSH layout). + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, + node_unit.Index(), + node_name, + 2u, + QNN_OP_CONCAT_PARAM_AXIS, + param_names)); + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_CONCAT, + {past_name, cur_name}, + {out_name}, + std::move(param_names), + do_op_validation), + "Failed to create KV concat node."); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// Helper: register an existing intermediate tensor as an APP_READ output +// by routing it through a no-op Reshape. +// --------------------------------------------------------------------------- +static Ort::Status RegisterIntermediateAsOutput(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const std::string& src_name, + const std::string& out_name, + const std::vector& shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool do_op_validation) { + QnnTensorWrapper out_tensor(out_name, QNN_TENSOR_TYPE_APP_READ, dtype, quant_param.Copy(), + std::vector(shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), + ("Failed to add output tensor: " + out_name).c_str()); + const std::string node_name = utils::UniqueNameGenerator().New(node_unit, "_out_reshape"); + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_RESHAPE, + {src_name}, + {out_name}, + {}, + do_op_validation), + "Failed to create output identity Reshape node."); + return Ort::Status(); +} + +// --------------------------------------------------------------------------- +// ProcessAttributesAndOutputs — emit the full decomposed attention graph +// --------------------------------------------------------------------------- +Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + std::vector&& input_names, + const Ort::Logger& /*logger*/, + bool do_op_validation) const { + const auto& onnx_inputs = node_unit.Inputs(); + const auto& onnx_outputs = node_unit.Outputs(); + + OrtNodeAttrHelper node_helper(node_unit); + + // ---- Gather input tensor info ---- + TensorInfo q_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[0], q_info)); + TensorInfo k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[1], k_info)); + TensorInfo v_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[2], v_info)); + + const Qnn_DataType_t dtype = q_info.qnn_data_type; + const QnnQuantParamsWrapper& q_quant = q_info.quant_param; + const QnnQuantParamsWrapper& k_quant = k_info.quant_param; + const QnnQuantParamsWrapper& v_quant = v_info.quant_param; + + const size_t q_rank = q_info.shape.size(); + const bool is_4d = (q_rank == 4); + + // ---- Resolve head counts and sequence/head dimensions ---- + // 3D: Q=[B,S_q,n_q*hs] K=[B,S_k,n_kv*hs] V=[B,S_k,n_kv*v_hs] + // 4D: Q=[B,n_q,S_q,hs] K=[B,n_kv,S_k,hs] V=[B,n_kv,S_k,v_hs] + const uint32_t B = q_info.shape[0]; + uint32_t n_q = 0, n_kv = 0, S_q = 0, S_k = 0, hs = 0, v_hs = 0; + + if (is_4d) { + n_q = q_info.shape[1]; + S_q = q_info.shape[2]; + hs = q_info.shape[3]; + n_kv = k_info.shape[1]; + S_k = k_info.shape[2]; + v_hs = v_info.shape[3]; + } else { + const auto opt_q_num_heads = node_helper.GetInt64("q_num_heads"); + const auto opt_kv_num_heads = node_helper.GetInt64("kv_num_heads"); + RETURN_IF_NOT(opt_q_num_heads.has_value() && opt_kv_num_heads.has_value(), + "q_num_heads and kv_num_heads are required for 3D Attention inputs"); + n_q = static_cast(opt_q_num_heads.value()); + n_kv = static_cast(opt_kv_num_heads.value()); + S_q = q_info.shape[1]; + S_k = k_info.shape[1]; + RETURN_IF(n_q == 0, "q_num_heads must be > 0"); + RETURN_IF(n_kv == 0, "kv_num_heads must be > 0"); + RETURN_IF(q_info.shape[2] % n_q != 0, "Q hidden dim must be divisible by q_num_heads"); + RETURN_IF(k_info.shape[2] % n_kv != 0, "K hidden dim must be divisible by kv_num_heads"); + hs = q_info.shape[2] / n_q; + v_hs = v_info.shape[2] / n_kv; + } + + // ---- Feature flags ---- + const bool is_gqa = (n_q != n_kv); + const uint32_t head_ratio = is_gqa ? (n_q / n_kv) : 1u; + + const bool has_past_key = (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()); + const bool has_qk_output = (onnx_outputs.size() > 3 && onnx_outputs[3].Exists()); + const int64_t qk_mode = node_helper.Get("qk_matmul_output_mode", static_cast(0)); + const float softcap = node_helper.Get("softcap", 0.0f); + + // ---- KV cache: resolve past seq dimension ---- + uint32_t S_past = 0; + if (has_past_key) { + TensorInfo past_k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[4], past_k_info)); + S_past = past_k_info.shape[2]; // [B, n, S_past, hs] + } + + // ---- Build input_names index helpers ---- + // The order in input_names matches the order ProcessInputs registered them: + // [0]=Q [1]=K [2]=V [3]=attn_mask (if present) [4]=past_key (if present) ... + const bool has_attn_mask = (onnx_inputs.size() > 3 && onnx_inputs[3].Exists()); + const std::string& q_in = input_names[0]; + const std::string& k_in = input_names[1]; + const std::string& v_in = input_names[2]; + // Offset into input_names for optional inputs (computed by walking present flags). + size_t opt_idx = 3; + const std::string attn_mask_in = has_attn_mask ? input_names[opt_idx++] : std::string{}; + const std::string past_key_in = has_past_key ? input_names[opt_idx++] : std::string{}; + const std::string past_value_in = has_past_key ? input_names[opt_idx++] : std::string{}; + + std::string q_cur = q_in; + std::string k_cur = k_in; + std::string v_cur = v_in; + + // ---- Step 1: Create a scalar initializer for sqrt(scale) ---- + const float scale_default = 1.0f / std::sqrt(static_cast(hs)); + const float scale_attr = node_helper.Get("scale", scale_default); + const float sqrt_scale = std::sqrt(scale_attr); + + const int64_t is_causal = node_helper.Get("is_causal", static_cast(0)); + + const std::string sqrt_scale_name = utils::UniqueNameGenerator().New(node_unit, "_sqrt_scale"); + { + std::vector scale_bytes; + if (dtype == QNN_DATATYPE_FLOAT_16) { + const Ort::Float16_t fp16_val(sqrt_scale); + scale_bytes.resize(sizeof(uint16_t)); + const uint16_t raw = fp16_val.val; + std::memcpy(scale_bytes.data(), &raw, sizeof(uint16_t)); + } else { + scale_bytes.resize(sizeof(float)); + std::memcpy(scale_bytes.data(), &sqrt_scale, sizeof(float)); + } + QnnTensorWrapper scale_tensor(sqrt_scale_name, + QNN_TENSOR_TYPE_STATIC, + dtype, + QnnQuantParamsWrapper{}, + {1u}, + std::move(scale_bytes)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(scale_tensor)), + "Failed to add sqrt_scale tensor."); + } + + // ---- Steps 2-3: Scale Q (always) and K (deferred when KV cache is active) ---- + // For KV cache: present_key = Concat(past_key_raw, K_raw) — store UNSCALED K. + // Scaling of the full K_present happens AFTER the concat so that past and + // current keys are treated uniformly. Without cache, scale K immediately. + const std::string q_scaled = utils::UniqueNameGenerator().New(node_unit, "_q_scaled"); + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + q_cur, sqrt_scale_name, q_scaled, + q_info.shape, dtype, q_quant, + /*is_graph_output=*/false, do_op_validation)); + q_cur = q_scaled; + + if (!has_past_key) { + // No KV cache: scale K now (standard path). + const std::string k_scaled = utils::UniqueNameGenerator().New(node_unit, "_k_scaled"); + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + k_cur, sqrt_scale_name, k_scaled, + k_info.shape, dtype, k_quant, + /*is_graph_output=*/false, do_op_validation)); + k_cur = k_scaled; + } + // If has_past_key: k_cur is still the raw (unscaled) K here. + // Scaling of k_present happens below, after the Concat. + + // ---- Steps 4-7 (3D only): Reshape + Transpose Q and K into [B, n, S, hs] ---- + if (!is_4d) { + // Reshape Q: [B, S_q, n_q*hs] -> [B, S_q, n_q, hs] + const std::string q_reshaped = utils::UniqueNameGenerator().New(node_unit, "_q_reshaped"); + const std::vector q_reshaped_shape = {B, S_q, n_q, hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(q_cur, q_reshaped, + {B, S_q, n_q * hs}, + q_reshaped_shape, + dtype, q_quant, + do_op_validation, + /*is_for_input=*/false)); + // Transpose Q: (0,2,1,3) -> [B, n_q, S_q, hs] + const std::string q_transposed = utils::UniqueNameGenerator().New(node_unit, "_q_transposed"); + const std::vector q_transposed_shape = {B, n_q, S_q, hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + q_reshaped, q_transposed, + q_reshaped_shape, + {0u, 2u, 1u, 3u}, + q_transposed_shape, + dtype, q_quant, + do_op_validation, + /*is_for_input=*/false)); + q_cur = q_transposed; + + // Reshape K: [B, S_k, n_kv*hs] -> [B, S_k, n_kv, hs] + const std::string k_reshaped = utils::UniqueNameGenerator().New(node_unit, "_k_reshaped"); + const std::vector k_reshaped_shape = {B, S_k, n_kv, hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(k_cur, k_reshaped, + {B, S_k, n_kv * hs}, + k_reshaped_shape, + dtype, k_quant, + do_op_validation, + /*is_for_input=*/false)); + // Transpose K: (0,2,1,3) -> [B, n_kv, S_k, hs] + const std::string k_transposed = utils::UniqueNameGenerator().New(node_unit, "_k_transposed"); + const std::vector k_transposed_shape = {B, n_kv, S_k, hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + k_reshaped, k_transposed, + k_reshaped_shape, + {0u, 2u, 1u, 3u}, + k_transposed_shape, + dtype, k_quant, + do_op_validation, + /*is_for_input=*/false)); + k_cur = k_transposed; + } + // k_cur is now [B, n_kv, S_k, hs] in BNSH layout. + + // ---- GQA expansion of K (if GQA) ---- + if (is_gqa) { + const std::vector k_in_shape = {B, n_kv, S_k, hs}; + const std::vector k_out_shape = {B, n_q, S_k, hs}; + const std::string k_expanded = utils::UniqueNameGenerator().New(node_unit, "_k_gqa"); + RETURN_IF_ERROR(AddGQAExpandNode(qnn_model_wrapper, node_unit, + k_cur, k_expanded, + k_in_shape, k_out_shape, + dtype, k_quant, + head_ratio, + do_op_validation)); + k_cur = k_expanded; + } + + // ---- KV cache concat for K ---- + std::string k_present_name; + if (has_past_key) { + const uint32_t S_k_total = S_past + S_k; + const std::vector k_present_shape = {B, n_q, S_k_total, hs}; + // If present_key is a graph output, emit it as APP_READ directly from concat. + const bool pk_is_graph_out = (onnx_outputs.size() > 1 && + onnx_outputs[1].Exists() && + qnn_model_wrapper.IsGraphOutput(onnx_outputs[1].name)); + k_present_name = pk_is_graph_out + ? onnx_outputs[1].name + : utils::UniqueNameGenerator().New(node_unit, "_k_present"); + RETURN_IF_ERROR(AddKVConcatNode(qnn_model_wrapper, node_unit, + past_key_in, k_cur, k_present_name, + k_present_shape, + dtype, k_quant, + pk_is_graph_out, + do_op_validation)); + k_cur = k_present_name; + S_k = S_k_total; // Update S_k to reflect the full sequence after concat. + + // Scale the full K_present (past + current) by sqrt(scale) for attention. + // This ensures past and current keys are treated uniformly. + const std::string k_present_scaled = utils::UniqueNameGenerator().New(node_unit, "_k_present_scaled"); + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + k_cur, sqrt_scale_name, k_present_scaled, + k_present_shape, dtype, k_quant, + /*is_graph_output=*/false, do_op_validation)); + k_cur = k_present_scaled; + } + + // ---- Steps 12-13 (3D only): Reshape + Transpose V into [B, n_kv, S_k_cur, v_hs] ---- + // NOTE: S_k used here is the original S_k of V (before KV concat). + // We do V transform BEFORE KV concat so V uses original S_k dimensions. + const uint32_t S_k_orig = has_past_key ? (S_k - S_past) : S_k; + std::string v_cur_4d = v_cur; + + if (!is_4d) { + // Reshape V: [B, S_k_orig, n_kv*v_hs] -> [B, S_k_orig, n_kv, v_hs] + const std::string v_reshaped = utils::UniqueNameGenerator().New(node_unit, "_v_reshaped"); + const std::vector v_reshaped_shape = {B, S_k_orig, n_kv, v_hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(v_cur, v_reshaped, + {B, S_k_orig, n_kv * v_hs}, + v_reshaped_shape, + dtype, v_quant, + do_op_validation, + /*is_for_input=*/false)); + // Transpose V: (0,2,1,3) -> [B, n_kv, S_k_orig, v_hs] + const std::string v_transposed = utils::UniqueNameGenerator().New(node_unit, "_v_transposed"); + const std::vector v_transposed_shape = {B, n_kv, S_k_orig, v_hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + v_reshaped, v_transposed, + v_reshaped_shape, + {0u, 2u, 1u, 3u}, + v_transposed_shape, + dtype, v_quant, + do_op_validation, + /*is_for_input=*/false)); + v_cur_4d = v_transposed; + } + // v_cur_4d is now [B, n_kv, S_k_orig, v_hs] in BNSH layout. + + // ---- GQA expansion of V (if GQA) ---- + if (is_gqa) { + const std::vector v_in_shape = {B, n_kv, S_k_orig, v_hs}; + const std::vector v_out_shape = {B, n_q, S_k_orig, v_hs}; + const std::string v_expanded = utils::UniqueNameGenerator().New(node_unit, "_v_gqa"); + RETURN_IF_ERROR(AddGQAExpandNode(qnn_model_wrapper, node_unit, + v_cur_4d, v_expanded, + v_in_shape, v_out_shape, + dtype, v_quant, + head_ratio, + do_op_validation)); + v_cur_4d = v_expanded; + } + + // ---- KV cache concat for V ---- + std::string v_present_name; + if (has_past_key) { + const uint32_t S_k_total = S_k; // already updated above. + const std::vector v_present_shape = {B, n_q, S_k_total, v_hs}; + const bool pv_is_graph_out = (onnx_outputs.size() > 2 && + onnx_outputs[2].Exists() && + qnn_model_wrapper.IsGraphOutput(onnx_outputs[2].name)); + v_present_name = pv_is_graph_out + ? onnx_outputs[2].name + : utils::UniqueNameGenerator().New(node_unit, "_v_present"); + RETURN_IF_ERROR(AddKVConcatNode(qnn_model_wrapper, node_unit, + past_value_in, v_cur_4d, v_present_name, + v_present_shape, + dtype, v_quant, + pv_is_graph_out, + do_op_validation)); + v_cur_4d = v_present_name; + } + + // ---- Step 8: MatMul Q * K^T (transpose_in1) -> [B, n_q, S_q, S_k] ---- + const std::string qk_out = utils::UniqueNameGenerator().New(node_unit, "_qk_out"); + const std::vector qk_shape = {B, n_q, S_q, S_k}; + RETURN_IF_ERROR(AddMatMulNode(qnn_model_wrapper, node_unit, + q_cur, k_cur, qk_out, + qk_shape, dtype, q_quant, + /*transpose_in1=*/true, do_op_validation)); + std::string scores_cur = qk_out; + + // ---- qk_matmul_output mode 0: capture post-QK scores ---- + std::string qk_captured; // The intermediate captured for qk_matmul_output. + if (has_qk_output && qk_mode == 0) { + qk_captured = scores_cur; + } + + // ---- Softcap ---- + if (softcap != 0.0f) { + const std::string sc_out = utils::UniqueNameGenerator().New(node_unit, "_scores_softcap"); + RETURN_IF_ERROR(AddSoftcapNode(qnn_model_wrapper, node_unit, + scores_cur, sc_out, + qk_shape, dtype, q_quant, + softcap, + do_op_validation)); + scores_cur = sc_out; + } + + // ---- qk_matmul_output mode 1: post-softcap scores ---- + if (has_qk_output && qk_mode == 1) { + qk_captured = scores_cur; + } + + // ---- Step 9 (is_causal=1): ADD static lower-triangular causal mask ---- + // With KV cache: offset = S_past so that row i attends to positions <= i+S_past. + if (is_causal != 0) { + const uint32_t offset = S_past; // 0 for no KV cache path. + const std::string causal_mask_name = utils::UniqueNameGenerator().New(node_unit, "_causal_mask"); + { + std::vector mask_shape = {B, n_q, S_q, S_k}; + const size_t total = static_cast(B) * static_cast(n_q) * + static_cast(S_q) * static_cast(S_k); + std::vector mask_bytes; + + if (dtype == QNN_DATATYPE_FLOAT_16) { + const Ort::Float16_t fp16_large_neg(-1e4f); + const uint16_t neg_raw = fp16_large_neg.val; + mask_bytes.resize(total * sizeof(uint16_t)); + uint16_t* mask_ptr = reinterpret_cast(mask_bytes.data()); + for (uint32_t b = 0; b < B; ++b) { + for (uint32_t h = 0; h < n_q; ++h) { + for (uint32_t i = 0; i < S_q; ++i) { + for (uint32_t j = 0; j < S_k; ++j) { + const size_t idx = ((static_cast(b) * n_q + h) * S_q + i) * S_k + j; + mask_ptr[idx] = (j <= i + offset) ? static_cast(0u) : neg_raw; + } + } + } + } + } else { + constexpr float large_neg = -1e9f; + mask_bytes.resize(total * sizeof(float)); + float* mask_ptr = reinterpret_cast(mask_bytes.data()); + for (uint32_t b = 0; b < B; ++b) { + for (uint32_t h = 0; h < n_q; ++h) { + for (uint32_t i = 0; i < S_q; ++i) { + for (uint32_t j = 0; j < S_k; ++j) { + const size_t idx = ((static_cast(b) * n_q + h) * S_q + i) * S_k + j; + mask_ptr[idx] = (j <= i + offset) ? 0.0f : large_neg; + } + } + } + } + } + + QnnTensorWrapper mask_tensor(causal_mask_name, + QNN_TENSOR_TYPE_STATIC, + dtype, + QnnQuantParamsWrapper{}, + std::move(mask_shape), + std::move(mask_bytes)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(mask_tensor)), + "Failed to add causal mask tensor."); + } + + const std::string masked_out = utils::UniqueNameGenerator().New(node_unit, "_causal_masked"); + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_ADD, + scores_cur, causal_mask_name, masked_out, + qk_shape, dtype, q_quant, + /*is_graph_output=*/false, do_op_validation)); + scores_cur = masked_out; + } + + // ---- Step 10 (attn_mask present): ADD user attention mask ---- + if (has_attn_mask) { + const std::string attn_masked_out = + utils::UniqueNameGenerator().New(node_unit, "_attn_masked"); + RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, + QNN_OP_ELEMENT_WISE_BINARY_OPERATION_ADD, + scores_cur, attn_mask_in, attn_masked_out, + qk_shape, dtype, q_quant, + /*is_graph_output=*/false, do_op_validation)); + scores_cur = attn_masked_out; + } + + // ---- qk_matmul_output mode 2: post-mask scores ---- + if (has_qk_output && qk_mode == 2) { + qk_captured = scores_cur; + } + + // ---- Step 11: Softmax (axis=3) -> [B, n_q, S_q, S_k] ---- + const std::string softmax_out = utils::UniqueNameGenerator().New(node_unit, "_softmax_out"); + RETURN_IF_ERROR(AddSoftmaxNode(qnn_model_wrapper, node_unit, + scores_cur, softmax_out, + qk_shape, dtype, q_quant, + /*axis=*/3u, do_op_validation)); + const std::string& attn_weights = softmax_out; + + // ---- qk_matmul_output mode 3: post-softmax (attn_weights) ---- + if (has_qk_output && qk_mode == 3) { + qk_captured = attn_weights; + } + + // ---- Step 14: MatMul attn_weights * V -> [B, n_q, S_q, v_hs] ---- + const std::string y_pre_transpose = + utils::UniqueNameGenerator().New(node_unit, "_y_pre_transpose"); + const std::vector y_pre_shape = {B, n_q, S_q, v_hs}; + RETURN_IF_ERROR(AddMatMulNode(qnn_model_wrapper, node_unit, + attn_weights, v_cur_4d, y_pre_transpose, + y_pre_shape, dtype, q_quant, + /*transpose_in1=*/false, do_op_validation)); + + // ---- Steps 15-16 (3D outputs): Transpose + Reshape back to [B, S_q, n_q*v_hs] ---- + const std::string& final_output_name = onnx_outputs[0].name; + const bool is_graph_output = qnn_model_wrapper.IsGraphOutput(final_output_name); + TensorInfo y_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[0], y_info)); + + if (!is_4d) { + // Transpose Y: (0,2,1,3) -> [B, S_q, n_q, v_hs] + const std::string y_transposed = utils::UniqueNameGenerator().New(node_unit, "_y_transposed"); + const std::vector y_transposed_shape = {B, S_q, n_q, v_hs}; + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + y_pre_transpose, y_transposed, + y_pre_shape, + {0u, 2u, 1u, 3u}, + y_transposed_shape, + dtype, q_quant, + do_op_validation, + /*is_for_input=*/false)); + // Reshape Y: [B, S_q, n_q, v_hs] -> [B, S_q, n_q*v_hs] + const std::vector y_final_shape = {B, S_q, n_q * v_hs}; + const Qnn_TensorType_t final_type = + is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + QnnTensorWrapper final_tensor(final_output_name, + final_type, + y_info.qnn_data_type, + y_info.quant_param.Copy(), + std::vector(y_final_shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(final_tensor)), + "Failed to add final Attention output tensor."); + + const std::string reshape_node = utils::UniqueNameGenerator().New(node_unit, "_y_reshape"); + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(reshape_node, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_RESHAPE, + {y_transposed}, + {final_output_name}, + {}, + do_op_validation), + "Failed to create final Reshape node."); + } else { + // 4D: y_pre_transpose is [B, n_q, S_q, v_hs] — already the correct output layout. + const Qnn_TensorType_t final_type = + is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + QnnTensorWrapper final_tensor(final_output_name, + final_type, + y_info.qnn_data_type, + y_info.quant_param.Copy(), + std::vector(y_pre_shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(final_tensor)), + "Failed to add final Attention output tensor (4D)."); + + const std::string rename_node = utils::UniqueNameGenerator().New(node_unit, "_y_rename"); + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(rename_node, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_RESHAPE, + {y_pre_transpose}, + {final_output_name}, + {}, + do_op_validation), + "Failed to create 4D output identity Reshape node."); + } + + // ---- Register KV cache outputs (if not already APP_READ) ---- + // When has_past_key is true and present_key/present_value are expected outputs, + // the concat nodes above already produced them. If they are graph outputs they + // were created as APP_READ. If they are not graph outputs but the ONNX node + // declares them as outputs we still need to register them. + if (has_past_key) { + // present_key = output[1] + if (onnx_outputs.size() > 1 && onnx_outputs[1].Exists()) { + const bool is_go = qnn_model_wrapper.IsGraphOutput(onnx_outputs[1].name); + if (!is_go && onnx_outputs[1].name != k_present_name) { + // Need to expose it. The concat output already has the right shape. + RETURN_IF_ERROR(RegisterIntermediateAsOutput(qnn_model_wrapper, node_unit, + k_present_name, + onnx_outputs[1].name, + {B, n_q, S_k, hs}, + dtype, k_quant, + do_op_validation)); + } + } + // present_value = output[2] + if (onnx_outputs.size() > 2 && onnx_outputs[2].Exists()) { + const bool is_go = qnn_model_wrapper.IsGraphOutput(onnx_outputs[2].name); + if (!is_go && onnx_outputs[2].name != v_present_name) { + RETURN_IF_ERROR(RegisterIntermediateAsOutput(qnn_model_wrapper, node_unit, + v_present_name, + onnx_outputs[2].name, + {B, n_q, S_k, v_hs}, + dtype, v_quant, + do_op_validation)); + } + } + } + + // ---- Register qk_matmul_output (output[3]) ---- + if (has_qk_output && !qk_captured.empty()) { + TensorInfo qk_out_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[3], qk_out_info)); + RETURN_IF_ERROR(RegisterIntermediateAsOutput(qnn_model_wrapper, node_unit, + qk_captured, + onnx_outputs[3].name, + qk_shape, + dtype, q_quant, + do_op_validation)); + } + + return Ort::Status(); +} + +void CreateAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + op_registrations.AddOpBuilder(op_type, std::make_unique()); +} + +#else // SDK version guard + +void CreateAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + ORT_UNUSED_PARAMETER(op_type); + ORT_UNUSED_PARAMETER(op_registrations); +} + +#endif // !(QNN_OPSET_VERSION_MAJOR < 2 || (QNN_OPSET_VERSION_MAJOR == 2 && QNN_OPSET_VERSION_MINOR <= 11)) + +} // namespace qnn +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc new file mode 100644 index 00000000000..193ce6079b9 --- /dev/null +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -0,0 +1,554 @@ +// Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. +// SPDX-License-Identifier: MIT + +#if !defined(ORT_MINIMAL_BUILD) + +#include +#include + +#include "gtest/gtest.h" + +#include "test/providers/qnn/qnn_test_utils.h" + +namespace onnxruntime { +namespace test { + +#if defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) + +// --------------------------------------------------------------------------- +// Helper: build a minimal float32 Attention test model (Q, K, V only). +// --------------------------------------------------------------------------- +static GetTestModelFn BuildAttentionTestCase( + const std::vector>& q_k_v_defs, + const std::vector& attrs) { + return [q_k_v_defs, attrs](ModelTestBuilder& builder) { + ASSERT_EQ(q_k_v_defs.size(), 3u); + + const std::string q_name = "attention_Q"; + const std::string k_name = "attention_K"; + const std::string v_name = "attention_V"; + + MakeTestInput(builder, q_name, q_k_v_defs[0]); + MakeTestInput(builder, k_name, q_k_v_defs[1]); + MakeTestInput(builder, v_name, q_k_v_defs[2]); + + builder.MakeOutput("attention_Y"); + builder.AddNode("attention_node", "Attention", + {q_name, k_name, v_name}, {"attention_Y"}, + kOnnxDomain, attrs); + }; +} + +// --------------------------------------------------------------------------- +// Helper: build an Attention model with an additive float attn_mask (input[3]). +// --------------------------------------------------------------------------- +static GetTestModelFn BuildAttentionTestCaseWithMask( + const std::vector>& q_k_v_defs, + const TestInputDef& mask_def, + const std::vector& attrs) { + return [q_k_v_defs, mask_def, attrs](ModelTestBuilder& builder) { + ASSERT_EQ(q_k_v_defs.size(), 3u); + + const std::string q_name = "attention_Q"; + const std::string k_name = "attention_K"; + const std::string v_name = "attention_V"; + const std::string mask_name = "attention_mask"; + + MakeTestInput(builder, q_name, q_k_v_defs[0]); + MakeTestInput(builder, k_name, q_k_v_defs[1]); + MakeTestInput(builder, v_name, q_k_v_defs[2]); + MakeTestInput(builder, mask_name, mask_def); + + builder.MakeOutput("attention_Y"); + builder.AddNode("attention_node", "Attention", + {q_name, k_name, v_name, mask_name}, {"attention_Y"}, + kOnnxDomain, attrs); + }; +} + +// --------------------------------------------------------------------------- +// Helper: Build an Attention model with KV cache (past_key/value as +// static initializers) and present_key/value as outputs. +// Q / K / V are dynamic inputs. +// past_key, past_value are static initializers (is_initializer=true). +// Outputs: Y, present_key, present_value. +// --------------------------------------------------------------------------- +static GetTestModelFn BuildAttentionTestCaseKV( + const TestInputDef& q_def, + const TestInputDef& k_def, + const TestInputDef& v_def, + const TestInputDef& past_key_def, + const TestInputDef& past_value_def, + const std::vector& attrs) { + return [q_def, k_def, v_def, past_key_def, past_value_def, attrs](ModelTestBuilder& builder) { + const std::string q_name = "attention_Q"; + const std::string k_name = "attention_K"; + const std::string v_name = "attention_V"; + const std::string past_key_name = "attention_past_key"; + const std::string past_val_name = "attention_past_value"; + + MakeTestInput(builder, q_name, q_def); + MakeTestInput(builder, k_name, k_def); + MakeTestInput(builder, v_name, v_def); + MakeTestInput(builder, past_key_name, past_key_def); + MakeTestInput(builder, past_val_name, past_value_def); + + builder.MakeOutput("attention_Y"); + builder.MakeOutput("attention_present_key"); + builder.MakeOutput("attention_present_value"); + + // ONNX Attention: inputs 0-5, outputs 0-2 (Y, present_key, present_value). + // Use an empty string "" for input[3] (no attn_mask) to skip that slot. + builder.AddNode("attention_node", "Attention", + {q_name, k_name, v_name, "", past_key_name, past_val_name}, + {"attention_Y", "attention_present_key", "attention_present_value"}, + kOnnxDomain, attrs); + }; +} + +// --------------------------------------------------------------------------- +// Helper: Build a 4D Attention model with qk_matmul_output (output[3]). +// --------------------------------------------------------------------------- +static GetTestModelFn BuildAttentionTestCaseDebugOutput( + const std::vector>& q_k_v_defs, + const std::vector& attrs) { + return [q_k_v_defs, attrs](ModelTestBuilder& builder) { + ASSERT_EQ(q_k_v_defs.size(), 3u); + + const std::string q_name = "attention_Q"; + const std::string k_name = "attention_K"; + const std::string v_name = "attention_V"; + + MakeTestInput(builder, q_name, q_k_v_defs[0]); + MakeTestInput(builder, k_name, q_k_v_defs[1]); + MakeTestInput(builder, v_name, q_k_v_defs[2]); + + builder.MakeOutput("attention_Y"); + // output[1] and output[2] (present_key/value) are absent. + // output[3] is qk_matmul_output. + builder.MakeOutput("attention_qk_output"); + + // Use empty strings for outputs[1] and outputs[2] to skip KV cache outputs. + builder.AddNode("attention_node", "Attention", + {q_name, k_name, v_name}, + {"attention_Y", "", "", "attention_qk_output"}, + kOnnxDomain, attrs); + }; +} + +// =========================================================================== +// 4D inputs (BNSH layout): no reshape/transpose needed +// =========================================================================== + +// MHA 4D non-causal. +// Q/K/V [1, 4, 8, 16]: batch=1, 4 heads, seq=8, head_size=16. +TEST_F(QnnHTPBackendTests, Attention_MHA_4D_NonCausal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 4D causal — adds static lower-triangular causal mask. +TEST_F(QnnHTPBackendTests, Attention_MHA_4D_Causal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 4D single head — n_heads=1 edge case. +TEST_F(QnnHTPBackendTests, Attention_MHA_4D_SingleHead) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 1, 8, 32}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 4D batch_size=2 — verifies batch dimension handling. +TEST_F(QnnHTPBackendTests, Attention_MHA_4D_Batch2) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({2, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({2, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({2, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 4D with explicit scale attribute. +TEST_F(QnnHTPBackendTests, Attention_MHA_4D_CustomScale) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("scale", 0.1f)}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// =========================================================================== +// 3D inputs (BSH layout): reshape + transpose to BNSH internally +// =========================================================================== + +// MHA 3D non-causal. +// Q/K/V [1, 8, 64]: seq=8, hidden=64=4*16, head_size=16. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_NonCausal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 3D causal — static lower-triangular mask. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_Causal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 3D single head. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_SingleHead) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 32}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(1)), + test::MakeAttribute("kv_num_heads", static_cast(1)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 3D batch_size=2. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_Batch2) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({2, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({2, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({2, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 3D with explicit additive attn_mask [1, 4, 8, 8]. +// Exercises the attn_mask ADD node in the decomposition. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_AttnMask) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + // Mask shape: [B=1, n_heads=4, S_q=8, S_k=8] — additive bias. + const std::vector mask_shape = {1, 4, 8, 8}; + + RunQnnModelTest( + BuildAttentionTestCaseWithMask( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + TestInputDef(mask_shape, false, -0.5f, 0.0f), // negative bias + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 3D causal + attn_mask — both mask paths active simultaneously. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_Causal_AttnMask) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + const std::vector mask_shape = {1, 4, 8, 8}; + + RunQnnModelTest( + BuildAttentionTestCaseWithMask( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + TestInputDef(mask_shape, false, -0.5f, 0.0f), + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MHA 3D with custom scale attribute. +TEST_F(QnnHTPBackendTests, Attention_MHA_3D_CustomScale) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("scale", 0.1f)}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// =========================================================================== +// GQA/MQA: q_num_heads != kv_num_heads +// =========================================================================== + +// GQA 3D non-causal — n_q=4, n_kv=2, head_ratio=2. +// Q[1,8,32] (4 heads * head_size=8), K/V[1,8,16] (2 heads * head_size=8). +TEST_F(QnnHTPBackendTests, Attention_GQA_3D_NonCausal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q: n_q=4, hs=8 + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // K: n_kv=2, hs=8 + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, // V: n_kv=2, hs=8 + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// GQA 3D causal with causal mask. +TEST_F(QnnHTPBackendTests, Attention_GQA_3D_Causal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// GQA 4D — Q[1,4,8,16], K/V[1,2,8,16], head_ratio=2. +TEST_F(QnnHTPBackendTests, Attention_GQA_4D) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // Q: [B,n_q,S,hs] + TestInputDef({1, 2, 8, 16}, false, -1.0f, 1.0f), // K: [B,n_kv,S,hs] + TestInputDef({1, 2, 8, 16}, false, -1.0f, 1.0f)}, // V: [B,n_kv,S,hs] + {test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// MQA 3D — n_q=4, n_kv=1 (Multi-Query Attention). +// Q[1,8,32] (4 heads * 8), K/V[1,8,8] (1 head * 8). +TEST_F(QnnHTPBackendTests, Attention_MQA_3D) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q: n_q=4, hs=8 + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f), // K: n_kv=1, hs=8 + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f)}, // V: n_kv=1, hs=8 + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(1)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// =========================================================================== +// Softcap: scores = softcap * tanh(scores / softcap) +// =========================================================================== + +// 4D MHA with softcap=10.0, non-causal. +TEST_F(QnnHTPBackendTests, Attention_Softcap_4D) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("softcap", 10.0f)}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// 3D MHA with softcap=50.0, causal. +// Q/K/V [1, 8, 64]: n_heads=4, head_size=16. +TEST_F(QnnHTPBackendTests, Attention_Softcap_3D_Causal) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(1)), + test::MakeAttribute("softcap", 50.0f)}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// =========================================================================== +// KV cache: past_key/value as static initializers, +// present_key/value as graph outputs. +// =========================================================================== + +// 4D MHA with KV cache. +// Q [1,4,8,16], K [1,4,8,16], V [1,4,8,16] +// past_key=[1,4,4,16] (S_past=4 initializer), past_value=[1,4,4,16] +// present_key=[1,4,12,16], present_value=[1,4,12,16] +TEST_F(QnnHTPBackendTests, Attention_KVCache_4D) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCaseKV( + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // Q + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // K + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // V + TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_key (initializer) + TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_value (initializer) + {test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + // KV cache (Concat) adds extra operations; fp16 rounding accumulates more error. + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(5e-2f)}); +} + +// =========================================================================== +// qk_matmul_output: debug outputs at different stages +// =========================================================================== + +// qk_matmul_output_mode=0 (post-QK matmul, before softcap/mask/softmax). +// Q/K/V [1,4,8,16]: qk_output shape should be [1,4,8,8]. +TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode0) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCaseDebugOutput( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("qk_matmul_output_mode", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// qk_matmul_output_mode=3 (post-softmax / attn_weights). +// qk_output shape should be [1,4,8,8] (same as mode 0 but different values). +TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode3) { + ProviderOptions opts; + opts["backend_type"] = "htp"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCaseDebugOutput( + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("qk_matmul_output_mode", static_cast(3))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +#endif // defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) + +} // namespace test +} // namespace onnxruntime + +#endif // !defined(ORT_MINIMAL_BUILD) From 5f0ef1838eac8043ddbb70283250bf5cdbfcf0b1 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Fri, 17 Jul 2026 12:24:32 -0700 Subject: [PATCH 03/20] Update documentation --- docs/execution_providers/QNN-ExecutionProvider.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/execution_providers/QNN-ExecutionProvider.md b/docs/execution_providers/QNN-ExecutionProvider.md index 408e1b0da65..8498db0585b 100644 --- a/docs/execution_providers/QNN-ExecutionProvider.md +++ b/docs/execution_providers/QNN-ExecutionProvider.md @@ -411,6 +411,7 @@ ort.unregister_execution_provider_library(ep_registration_name) |ai.onnx:ArgMin|| |ai.onnx:Asin|| |ai.onnx:Atan|| +|ai.onnx:Attention|Opsets 23–24. 3D \[B,S,n·hs\] and 4D \[B,n,S,hs\] inputs. MHA and GQA/MQA (q\_num\_heads ≥ kv\_num\_heads, divisible). is\_causal, attn\_mask, softcap, KV cache (past/present key·value), qk\_matmul\_output modes 0–3. Static shapes only.| |ai.onnx:AveragePool|| |ai.onnx:BatchNormalization|fp16 supported since 1.18.0| |ai.onnx:Cast|| From e8b97ce6c67a9a4534b9459b55e701ea4fbc1ee8 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Fri, 17 Jul 2026 12:26:38 -0700 Subject: [PATCH 04/20] Lint fixes --- .../builder/opbuilder/attention_op_builder.cc | 210 +++++++++--------- .../test/providers/qnn/attention_test.cc | 110 ++++----- 2 files changed, 160 insertions(+), 160 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 9255b25ceb3..8aec23bf18c 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -54,8 +54,8 @@ class AttentionOpBuilder : public BaseOpBuilder { // IsOpSupported // --------------------------------------------------------------------------- Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const Ort::Logger& logger) const { + const OrtNodeUnit& node_unit, + const Ort::Logger& logger) const { ORT_UNUSED_PARAMETER(logger); const auto& inputs = node_unit.Inputs(); @@ -177,10 +177,10 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper // nonpad_kv_seqlen // --------------------------------------------------------------------------- Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const Ort::Logger& logger, - std::vector& input_names, - bool /*do_op_validation*/) const { + const OrtNodeUnit& node_unit, + const Ort::Logger& logger, + std::vector& input_names, + bool /*do_op_validation*/) const { const auto& onnx_inputs = node_unit.Inputs(); RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[0], logger, input_names)); // Q @@ -211,16 +211,16 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper // Helper: emit an ElementWiseBinary (MUL or ADD or DIV) node. // --------------------------------------------------------------------------- static Ort::Status AddBinaryOpNode(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - uint32_t operation, - const std::string& lhs_name, - const std::string& rhs_name, - const std::string& out_name, - const std::vector& out_shape, - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - bool is_graph_output, - bool do_op_validation) { + const OrtNodeUnit& node_unit, + uint32_t operation, + const std::string& lhs_name, + const std::string& rhs_name, + const std::string& out_name, + const std::vector& out_shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool is_graph_output, + bool do_op_validation) { const Qnn_TensorType_t tensor_type = is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; QnnTensorWrapper out_tensor(out_name, tensor_type, dtype, quant_param.Copy(), @@ -232,11 +232,11 @@ static Ort::Status AddBinaryOpNode(QnnModelWrapper& qnn_model_wrapper, std::vector param_names; RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, - node_unit.Index(), - node_name, - operation, - QNN_OP_ELEMENT_WISE_BINARY_PARAM_OPERATION, - param_names)); + node_unit.Index(), + node_name, + operation, + QNN_OP_ELEMENT_WISE_BINARY_PARAM_OPERATION, + param_names)); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, QNN_OP_PACKAGE_NAME_QTI_AISW, @@ -253,15 +253,15 @@ static Ort::Status AddBinaryOpNode(QnnModelWrapper& qnn_model_wrapper, // Helper: emit a MatMul node (with optional transpose_in1). // --------------------------------------------------------------------------- static Ort::Status AddMatMulNode(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const std::string& lhs_name, - const std::string& rhs_name, - const std::string& out_name, - const std::vector& out_shape, - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - bool transpose_in1, - bool do_op_validation) { + const OrtNodeUnit& node_unit, + const std::string& lhs_name, + const std::string& rhs_name, + const std::string& out_name, + const std::vector& out_shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool transpose_in1, + bool do_op_validation) { QnnTensorWrapper out_tensor(out_name, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), std::vector(out_shape)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), @@ -294,14 +294,14 @@ static Ort::Status AddMatMulNode(QnnModelWrapper& qnn_model_wrapper, // Helper: emit a Softmax node (axis param). // --------------------------------------------------------------------------- static Ort::Status AddSoftmaxNode(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const std::string& in_name, - const std::string& out_name, - const std::vector& shape, - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - uint32_t axis, - bool do_op_validation) { + const OrtNodeUnit& node_unit, + const std::string& in_name, + const std::string& out_name, + const std::vector& shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + uint32_t axis, + bool do_op_validation) { QnnTensorWrapper out_tensor(out_name, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), std::vector(shape)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), @@ -311,11 +311,11 @@ static Ort::Status AddSoftmaxNode(QnnModelWrapper& qnn_model_wrapper, std::vector param_names; RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, - node_unit.Index(), - node_name, - axis, - QNN_OP_SOFTMAX_PARAM_AXIS, - param_names)); + node_unit.Index(), + node_name, + axis, + QNN_OP_SOFTMAX_PARAM_AXIS, + param_names)); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, QNN_OP_PACKAGE_NAME_QTI_AISW, @@ -332,9 +332,9 @@ static Ort::Status AddSoftmaxNode(QnnModelWrapper& qnn_model_wrapper, // Helper: create a static scalar tensor (fp32 or fp16) for softcap arithmetic. // --------------------------------------------------------------------------- static Ort::Status AddScalarTensor(QnnModelWrapper& qnn_model_wrapper, - const std::string& name, - float value, - Qnn_DataType_t dtype) { + const std::string& name, + float value, + Qnn_DataType_t dtype) { std::vector bytes; if (dtype == QNN_DATATYPE_FLOAT_16) { const Ort::Float16_t fp16(value); @@ -358,14 +358,14 @@ static Ort::Status AddScalarTensor(QnnModelWrapper& qnn_model_wrapper, // Steps: Div(scores, sc) → Tanh → Mul(result, sc) // --------------------------------------------------------------------------- static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const std::string& in_name, - const std::string& out_name, - const std::vector& shape, - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - float softcap_val, - bool do_op_validation) { + const OrtNodeUnit& node_unit, + const std::string& in_name, + const std::string& out_name, + const std::vector& shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + float softcap_val, + bool do_op_validation) { // Static scalar for softcap value. const std::string sc_name = utils::UniqueNameGenerator().New(node_unit, "_softcap_scalar"); RETURN_IF_ERROR(AddScalarTensor(qnn_model_wrapper, sc_name, softcap_val, dtype)); @@ -417,15 +417,15 @@ static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, // // --------------------------------------------------------------------------- static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const std::string& in_name, - const std::string& out_name, - const std::vector& in_shape, // [B, n_kv, S, hs] - const std::vector& out_shape, // [B, n_q, S, hs] - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - uint32_t head_ratio, - bool /*do_op_validation*/) { + const OrtNodeUnit& node_unit, + const std::string& in_name, + const std::string& out_name, + const std::vector& in_shape, // [B, n_kv, S, hs] + const std::vector& out_shape, // [B, n_q, S, hs] + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + uint32_t head_ratio, + bool /*do_op_validation*/) { // 4D-only GQA expansion that avoids 5D tensors (HTP finalization fails with 5D). // // Goal: produce K_expanded[b, kv*head_ratio+r, s, h] = K[b, kv, s, h] @@ -439,11 +439,11 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, // → Transpose (0,2,1,3) → [B, n_kv, head_ratio, S*hs] // → Reshape [B, n_q, S, hs] (C-order: kv*head_ratio+r → floor-div ✓) - const uint32_t B = in_shape[0]; + const uint32_t B = in_shape[0]; const uint32_t n_kv = in_shape[1]; - const uint32_t S = in_shape[2]; - const uint32_t hs = in_shape[3]; - const uint32_t Shs = S * hs; // merged dim + const uint32_t S = in_shape[2]; + const uint32_t hs = in_shape[3]; + const uint32_t Shs = S * hs; // merged dim // Step 1: Reshape [B, n_kv, S, hs] → [B, 1, n_kv, S*hs] const std::string r1_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_r1"); @@ -484,7 +484,7 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, } // Step 3: Transpose (0,2,1,3) → [B, n_kv, head_ratio, S*hs] - const std::string tr_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_tr"); + const std::string tr_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_tr"); const std::vector tr_shape = {B, n_kv, head_ratio, Shs}; RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), tiled_name, tr_name, @@ -511,15 +511,15 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, // out_shape = [B, n, S_past+S_cur, hs] // --------------------------------------------------------------------------- static Ort::Status AddKVConcatNode(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const std::string& past_name, - const std::string& cur_name, - const std::string& out_name, - const std::vector& out_shape, - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - bool is_graph_output, - bool do_op_validation) { + const OrtNodeUnit& node_unit, + const std::string& past_name, + const std::string& cur_name, + const std::string& out_name, + const std::vector& out_shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool is_graph_output, + bool do_op_validation) { const Qnn_TensorType_t tensor_type = is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; QnnTensorWrapper out_tensor(out_name, tensor_type, dtype, quant_param.Copy(), @@ -531,11 +531,11 @@ static Ort::Status AddKVConcatNode(QnnModelWrapper& qnn_model_wrapper, std::vector param_names; // axis = 2 (the sequence dimension in BNSH layout). RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, - node_unit.Index(), - node_name, - 2u, - QNN_OP_CONCAT_PARAM_AXIS, - param_names)); + node_unit.Index(), + node_name, + 2u, + QNN_OP_CONCAT_PARAM_AXIS, + param_names)); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, QNN_OP_PACKAGE_NAME_QTI_AISW, QNN_OP_CONCAT, @@ -552,13 +552,13 @@ static Ort::Status AddKVConcatNode(QnnModelWrapper& qnn_model_wrapper, // by routing it through a no-op Reshape. // --------------------------------------------------------------------------- static Ort::Status RegisterIntermediateAsOutput(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const std::string& src_name, - const std::string& out_name, - const std::vector& shape, - Qnn_DataType_t dtype, - const QnnQuantParamsWrapper& quant_param, - bool do_op_validation) { + const OrtNodeUnit& node_unit, + const std::string& src_name, + const std::string& out_name, + const std::vector& shape, + Qnn_DataType_t dtype, + const QnnQuantParamsWrapper& quant_param, + bool do_op_validation) { QnnTensorWrapper out_tensor(out_name, QNN_TENSOR_TYPE_APP_READ, dtype, quant_param.Copy(), std::vector(shape)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(out_tensor)), @@ -579,10 +579,10 @@ static Ort::Status RegisterIntermediateAsOutput(QnnModelWrapper& qnn_model_wrapp // ProcessAttributesAndOutputs — emit the full decomposed attention graph // --------------------------------------------------------------------------- Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - std::vector&& input_names, - const Ort::Logger& /*logger*/, - bool do_op_validation) const { + const OrtNodeUnit& node_unit, + std::vector&& input_names, + const Ort::Logger& /*logger*/, + bool do_op_validation) const { const auto& onnx_inputs = node_unit.Inputs(); const auto& onnx_outputs = node_unit.Outputs(); @@ -791,8 +791,8 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const std::vector k_present_shape = {B, n_q, S_k_total, hs}; // If present_key is a graph output, emit it as APP_READ directly from concat. const bool pk_is_graph_out = (onnx_outputs.size() > 1 && - onnx_outputs[1].Exists() && - qnn_model_wrapper.IsGraphOutput(onnx_outputs[1].name)); + onnx_outputs[1].Exists() && + qnn_model_wrapper.IsGraphOutput(onnx_outputs[1].name)); k_present_name = pk_is_graph_out ? onnx_outputs[1].name : utils::UniqueNameGenerator().New(node_unit, "_k_present"); @@ -867,8 +867,8 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const uint32_t S_k_total = S_k; // already updated above. const std::vector v_present_shape = {B, n_q, S_k_total, v_hs}; const bool pv_is_graph_out = (onnx_outputs.size() > 2 && - onnx_outputs[2].Exists() && - qnn_model_wrapper.IsGraphOutput(onnx_outputs[2].name)); + onnx_outputs[2].Exists() && + qnn_model_wrapper.IsGraphOutput(onnx_outputs[2].name)); v_present_name = pv_is_graph_out ? onnx_outputs[2].name : utils::UniqueNameGenerator().New(node_unit, "_v_present"); @@ -993,9 +993,9 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // ---- Step 11: Softmax (axis=3) -> [B, n_q, S_q, S_k] ---- const std::string softmax_out = utils::UniqueNameGenerator().New(node_unit, "_softmax_out"); RETURN_IF_ERROR(AddSoftmaxNode(qnn_model_wrapper, node_unit, - scores_cur, softmax_out, - qk_shape, dtype, q_quant, - /*axis=*/3u, do_op_validation)); + scores_cur, softmax_out, + qk_shape, dtype, q_quant, + /*axis=*/3u, do_op_validation)); const std::string& attn_weights = softmax_out; // ---- qk_matmul_output mode 3: post-softmax (attn_weights) ---- @@ -1112,11 +1112,11 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn TensorInfo qk_out_info{}; RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[3], qk_out_info)); RETURN_IF_ERROR(RegisterIntermediateAsOutput(qnn_model_wrapper, node_unit, - qk_captured, - onnx_outputs[3].name, - qk_shape, - dtype, q_quant, - do_op_validation)); + qk_captured, + onnx_outputs[3].name, + qk_shape, + dtype, q_quant, + do_op_validation)); } return Ort::Status(); diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 193ce6079b9..266a0039934 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -49,14 +49,14 @@ static GetTestModelFn BuildAttentionTestCaseWithMask( return [q_k_v_defs, mask_def, attrs](ModelTestBuilder& builder) { ASSERT_EQ(q_k_v_defs.size(), 3u); - const std::string q_name = "attention_Q"; - const std::string k_name = "attention_K"; - const std::string v_name = "attention_V"; + const std::string q_name = "attention_Q"; + const std::string k_name = "attention_K"; + const std::string v_name = "attention_V"; const std::string mask_name = "attention_mask"; - MakeTestInput(builder, q_name, q_k_v_defs[0]); - MakeTestInput(builder, k_name, q_k_v_defs[1]); - MakeTestInput(builder, v_name, q_k_v_defs[2]); + MakeTestInput(builder, q_name, q_k_v_defs[0]); + MakeTestInput(builder, k_name, q_k_v_defs[1]); + MakeTestInput(builder, v_name, q_k_v_defs[2]); MakeTestInput(builder, mask_name, mask_def); builder.MakeOutput("attention_Y"); @@ -81,15 +81,15 @@ static GetTestModelFn BuildAttentionTestCaseKV( const TestInputDef& past_value_def, const std::vector& attrs) { return [q_def, k_def, v_def, past_key_def, past_value_def, attrs](ModelTestBuilder& builder) { - const std::string q_name = "attention_Q"; - const std::string k_name = "attention_K"; - const std::string v_name = "attention_V"; - const std::string past_key_name = "attention_past_key"; - const std::string past_val_name = "attention_past_value"; - - MakeTestInput(builder, q_name, q_def); - MakeTestInput(builder, k_name, k_def); - MakeTestInput(builder, v_name, v_def); + const std::string q_name = "attention_Q"; + const std::string k_name = "attention_K"; + const std::string v_name = "attention_V"; + const std::string past_key_name = "attention_past_key"; + const std::string past_val_name = "attention_past_value"; + + MakeTestInput(builder, q_name, q_def); + MakeTestInput(builder, k_name, k_def); + MakeTestInput(builder, v_name, v_def); MakeTestInput(builder, past_key_name, past_key_def); MakeTestInput(builder, past_val_name, past_value_def); @@ -217,7 +217,7 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_4D_CustomScale) { TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, {test::MakeAttribute("is_causal", static_cast(0)), - test::MakeAttribute("scale", 0.1f)}), + test::MakeAttribute("scale", 0.1f)}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -238,9 +238,9 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_NonCausal) { {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -256,9 +256,9 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_Causal) { {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(1))}), + test::MakeAttribute("is_causal", static_cast(1))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -274,9 +274,9 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_SingleHead) { {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), TestInputDef({1, 8, 32}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(1)), + {test::MakeAttribute("q_num_heads", static_cast(1)), test::MakeAttribute("kv_num_heads", static_cast(1)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -292,9 +292,9 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_Batch2) { {TestInputDef({2, 8, 64}, false, -1.0f, 1.0f), TestInputDef({2, 8, 64}, false, -1.0f, 1.0f), TestInputDef({2, 8, 64}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -315,9 +315,9 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_AttnMask) { TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, TestInputDef(mask_shape, false, -0.5f, 0.0f), // negative bias - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -336,9 +336,9 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_Causal_AttnMask) { TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, TestInputDef(mask_shape, false, -0.5f, 0.0f), - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(1))}), + test::MakeAttribute("is_causal", static_cast(1))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -354,10 +354,10 @@ TEST_F(QnnHTPBackendTests, Attention_MHA_3D_CustomScale) { {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(0)), - test::MakeAttribute("scale", 0.1f)}), + test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("scale", 0.1f)}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -378,9 +378,9 @@ TEST_F(QnnHTPBackendTests, Attention_GQA_3D_NonCausal) { {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q: n_q=4, hs=8 TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // K: n_kv=2, hs=8 TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, // V: n_kv=2, hs=8 - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(2)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -396,9 +396,9 @@ TEST_F(QnnHTPBackendTests, Attention_GQA_3D_Causal) { {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(2)), - test::MakeAttribute("is_causal", static_cast(1))}), + test::MakeAttribute("is_causal", static_cast(1))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -411,9 +411,9 @@ TEST_F(QnnHTPBackendTests, Attention_GQA_4D) { RunQnnModelTest( BuildAttentionTestCase( - {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // Q: [B,n_q,S,hs] - TestInputDef({1, 2, 8, 16}, false, -1.0f, 1.0f), // K: [B,n_kv,S,hs] - TestInputDef({1, 2, 8, 16}, false, -1.0f, 1.0f)}, // V: [B,n_kv,S,hs] + {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // Q: [B,n_q,S,hs] + TestInputDef({1, 2, 8, 16}, false, -1.0f, 1.0f), // K: [B,n_kv,S,hs] + TestInputDef({1, 2, 8, 16}, false, -1.0f, 1.0f)}, // V: [B,n_kv,S,hs] {test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); @@ -429,11 +429,11 @@ TEST_F(QnnHTPBackendTests, Attention_MQA_3D) { RunQnnModelTest( BuildAttentionTestCase( {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q: n_q=4, hs=8 - TestInputDef({1, 8, 8}, false, -1.0f, 1.0f), // K: n_kv=1, hs=8 - TestInputDef({1, 8, 8}, false, -1.0f, 1.0f)}, // V: n_kv=1, hs=8 - {test::MakeAttribute("q_num_heads", static_cast(4)), + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f), // K: n_kv=1, hs=8 + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f)}, // V: n_kv=1, hs=8 + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(1)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -454,7 +454,7 @@ TEST_F(QnnHTPBackendTests, Attention_Softcap_4D) { TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, {test::MakeAttribute("is_causal", static_cast(0)), - test::MakeAttribute("softcap", 10.0f)}), + test::MakeAttribute("softcap", 10.0f)}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -471,10 +471,10 @@ TEST_F(QnnHTPBackendTests, Attention_Softcap_3D_Causal) { {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), + {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(1)), - test::MakeAttribute("softcap", 50.0f)}), + test::MakeAttribute("is_causal", static_cast(1)), + test::MakeAttribute("softcap", 50.0f)}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -495,11 +495,11 @@ TEST_F(QnnHTPBackendTests, Attention_KVCache_4D) { RunQnnModelTest( BuildAttentionTestCaseKV( - TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // Q - TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // K - TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // V - TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_key (initializer) - TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_value (initializer) + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // Q + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // K + TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), // V + TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_key (initializer) + TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_value (initializer) {test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, // KV cache (Concat) adds extra operations; fp16 rounding accumulates more error. @@ -522,8 +522,8 @@ TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode0) { {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("is_causal", static_cast(0)), - test::MakeAttribute("qk_matmul_output_mode", static_cast(0))}), + {test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("qk_matmul_output_mode", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -540,8 +540,8 @@ TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode3) { {TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f), TestInputDef({1, 4, 8, 16}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("is_causal", static_cast(0)), - test::MakeAttribute("qk_matmul_output_mode", static_cast(3))}), + {test::MakeAttribute("is_causal", static_cast(0)), + test::MakeAttribute("qk_matmul_output_mode", static_cast(3))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } From 2d6a7d3a8addcf23420f437a949ba5ecb0983449 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Sun, 19 Jul 2026 17:44:55 -0700 Subject: [PATCH 05/20] Update tolerance to 2e-3f --- onnxruntime/test/providers/qnn/attention_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 266a0039934..f18d4d2bdf7 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -503,7 +503,7 @@ TEST_F(QnnHTPBackendTests, Attention_KVCache_4D) { {test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, // KV cache (Concat) adds extra operations; fp16 rounding accumulates more error. - EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(5e-2f)}); + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } // =========================================================================== From 2093cf0763a8762790808c61e59fe64546571723 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Sun, 19 Jul 2026 17:59:58 -0700 Subject: [PATCH 06/20] Decomposition fix --- .../builder/opbuilder/attention_op_builder.cc | 64 ++++++++++++------- 1 file changed, 41 insertions(+), 23 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 8aec23bf18c..184b20aa5cb 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -378,7 +378,7 @@ static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, shape, dtype, quant_param, /*is_graph_output=*/false, do_op_validation)); - // Tanh(x) -> t + // Tanh(x) -> t [QNN_OP_ELEMENT_WISE_NEURON with TANH param] const std::string tanh_out = utils::UniqueNameGenerator().New(node_unit, "_softcap_tanh"); { QnnTensorWrapper tanh_tensor(tanh_out, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), @@ -386,12 +386,22 @@ static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(tanh_tensor)), ("Failed to add softcap Tanh output: " + tanh_out).c_str()); const std::string tanh_node = utils::UniqueNameGenerator().New(node_unit, "_tanh"); + + Qnn_Scalar_t tanh_op_scalar = QNN_SCALAR_INIT; + tanh_op_scalar.dataType = QNN_DATATYPE_UINT_32; + tanh_op_scalar.uint32Value = QNN_OP_ELEMENT_WISE_NEURON_OPERATION_TANH; + QnnParamWrapper tanh_op_param(node_unit.Index(), tanh_node, + QNN_OP_ELEMENT_WISE_NEURON_PARAM_OPERATION, + tanh_op_scalar); + std::vector tanh_param_names = {tanh_op_param.GetParamTensorName()}; + RETURN_IF_NOT(qnn_model_wrapper.AddParamWrapper(std::move(tanh_op_param)), + "Failed to add softcap Tanh operation param."); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(tanh_node, QNN_OP_PACKAGE_NAME_QTI_AISW, - QNN_OP_TANH, + QNN_OP_ELEMENT_WISE_NEURON, {div_out}, {tanh_out}, - {}, + std::move(tanh_param_names), do_op_validation), "Failed to create softcap Tanh node."); } @@ -581,7 +591,7 @@ static Ort::Status RegisterIntermediateAsOutput(QnnModelWrapper& qnn_model_wrapp Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, std::vector&& input_names, - const Ort::Logger& /*logger*/, + const Ort::Logger& logger, bool do_op_validation) const { const auto& onnx_inputs = node_unit.Inputs(); const auto& onnx_outputs = node_unit.Outputs(); @@ -890,29 +900,21 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn /*transpose_in1=*/true, do_op_validation)); std::string scores_cur = qk_out; - // ---- qk_matmul_output mode 0: capture post-QK scores ---- + // ---- qk_matmul_output mode 0: raw post-QK scores (pre-mask, pre-softcap) ---- + // Per ONNX spec the pipeline is: QK → masks → softcap → softmax + // mode 0 = raw QK, mode 1 = post-mask pre-softcap, + // mode 2 = post-softcap, mode 3 = post-softmax. std::string qk_captured; // The intermediate captured for qk_matmul_output. if (has_qk_output && qk_mode == 0) { qk_captured = scores_cur; } - // ---- Softcap ---- - if (softcap != 0.0f) { - const std::string sc_out = utils::UniqueNameGenerator().New(node_unit, "_scores_softcap"); - RETURN_IF_ERROR(AddSoftcapNode(qnn_model_wrapper, node_unit, - scores_cur, sc_out, - qk_shape, dtype, q_quant, - softcap, - do_op_validation)); - scores_cur = sc_out; - } - - // ---- qk_matmul_output mode 1: post-softcap scores ---- - if (has_qk_output && qk_mode == 1) { - qk_captured = scores_cur; - } + // ---- Masks applied BEFORE softcap (per ONNX spec) ---- + // Applying mask after softcap would be wrong: softcap saturates large negatives + // to -softcap (not -inf), so the mask must set large negatives first so softcap + // can compress the full range uniformly. - // ---- Step 9 (is_causal=1): ADD static lower-triangular causal mask ---- + // ---- Causal mask (is_causal=1): ADD static lower-triangular mask ---- // With KV cache: offset = S_past so that row i attends to positions <= i+S_past. if (is_causal != 0) { const uint32_t offset = S_past; // 0 for no KV cache path. @@ -973,7 +975,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn scores_cur = masked_out; } - // ---- Step 10 (attn_mask present): ADD user attention mask ---- + // ---- User attention mask: ADD ---- if (has_attn_mask) { const std::string attn_masked_out = utils::UniqueNameGenerator().New(node_unit, "_attn_masked"); @@ -985,7 +987,23 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn scores_cur = attn_masked_out; } - // ---- qk_matmul_output mode 2: post-mask scores ---- + // ---- qk_matmul_output mode 1: post-mask, pre-softcap ---- + if (has_qk_output && qk_mode == 1) { + qk_captured = scores_cur; + } + + // ---- Softcap (applied AFTER masks per ONNX spec) ---- + if (softcap != 0.0f) { + const std::string sc_out = utils::UniqueNameGenerator().New(node_unit, "_scores_softcap"); + RETURN_IF_ERROR(AddSoftcapNode(qnn_model_wrapper, node_unit, + scores_cur, sc_out, + qk_shape, dtype, q_quant, + softcap, + do_op_validation)); + scores_cur = sc_out; + } + + // ---- qk_matmul_output mode 2: post-softcap ---- if (has_qk_output && qk_mode == 2) { qk_captured = scores_cur; } From 734f7015074c8f64f300d2442978d3e3ac505841 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Sun, 19 Jul 2026 19:28:39 -0700 Subject: [PATCH 07/20] Build fix --- .../providers/qnn/builder/opbuilder/attention_op_builder.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 184b20aa5cb..21c95078961 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -591,7 +591,7 @@ static Ort::Status RegisterIntermediateAsOutput(QnnModelWrapper& qnn_model_wrapp Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, std::vector&& input_names, - const Ort::Logger& logger, + const Ort::Logger& /*logger*/, bool do_op_validation) const { const auto& onnx_inputs = node_unit.Inputs(); const auto& onnx_outputs = node_unit.Outputs(); From 24fc679df5b0e9808aff0f8a4d6fe1083bddd77f Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Sun, 19 Jul 2026 20:51:55 -0700 Subject: [PATCH 08/20] Gate Softcap tests to ARM64 --- onnxruntime/test/providers/qnn/attention_test.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index f18d4d2bdf7..a2424db4ca1 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -440,7 +440,11 @@ TEST_F(QnnHTPBackendTests, Attention_MQA_3D) { // =========================================================================== // Softcap: scores = softcap * tanh(scores / softcap) +// Gated to real ARM64 hardware: the softcap chain (Div + ElementWiseNeuron +// TANH + Mul) triggers QNN_COMMON_ERROR_MEM_ALLOC during HTP graph +// finalization on the x86_64 HTP simulator. // =========================================================================== +#if defined(__aarch64__) || defined(_M_ARM64) // 4D MHA with softcap=10.0, non-causal. TEST_F(QnnHTPBackendTests, Attention_Softcap_4D) { @@ -479,6 +483,8 @@ TEST_F(QnnHTPBackendTests, Attention_Softcap_3D_Causal) { EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } +#endif // defined(__aarch64__) || defined(_M_ARM64) + // =========================================================================== // KV cache: past_key/value as static initializers, // present_key/value as graph outputs. From 5d91252fd24b3e91b3d7072cf4a9e728dfa46567 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 23 Jul 2026 14:57:34 -0700 Subject: [PATCH 09/20] [QNN EP] Use native GQA op for GPU --- .../builder/opbuilder/attention_op_builder.cc | 244 ++++++++++++++++++ .../test/providers/qnn/attention_test.cc | 175 +++++++++++++ 2 files changed, 419 insertions(+) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 21c95078961..0380ff02018 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -183,6 +183,35 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper bool /*do_op_validation*/) const { const auto& onnx_inputs = node_unit.Inputs(); + // ---- GPU native GQA path: determine routing ---- + { + TensorInfo q_info{}, k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[0], q_info)); + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[1], k_info)); + const size_t q_rank = q_info.shape.size(); + OrtNodeAttrHelper node_helper(node_unit); + uint32_t n_q = 0, n_kv = 0; + if (q_rank == 4) { + n_q = q_info.shape[1]; + n_kv = k_info.shape[1]; + } else { + const auto opt_q = node_helper.GetInt64("q_num_heads"); + const auto opt_kv = node_helper.GetInt64("kv_num_heads"); + if (opt_q.has_value()) n_q = static_cast(opt_q.value()); + if (opt_kv.has_value()) n_kv = static_cast(opt_kv.value()); + } + const int64_t is_causal = node_helper.Get("is_causal", static_cast(0)); + const float softcap = node_helper.Get("softcap", 0.0f); + const bool has_attn_mask = (onnx_inputs.size() > 3 && onnx_inputs[3].Exists()); + const bool has_qk_output = (node_unit.Outputs().size() > 3 && + node_unit.Outputs()[3].Exists()); + if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), q_rank, + n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output)) { + return ProcessInputsNativeGQA(qnn_model_wrapper, node_unit, logger, input_names); + } + } + + // ---- Decomposition path ---- RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[0], logger, input_names)); // Q RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[1], logger, input_names)); // K RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[2], logger, input_names)); // V @@ -207,6 +236,213 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper return Ort::Status(); } +// --------------------------------------------------------------------------- +// GPU native GQA path +// +// When backend=GPU, input layout is 3D BSH, the node is true GQA +// (n_q > n_kv, divisible), is_causal=1, and no features QNN GQA cannot +// express (softcap/attn_mask/qk_output), a single +// QNN_OP_GROUP_QUERY_ATTENTION node is emitted instead of the decomposed +// graph. All other cases (MHA, non-causal, softcap, HTP, 4D inputs) fall +// through to decomposition unchanged. +// --------------------------------------------------------------------------- + +static bool ShouldUseNativeGQA(QnnBackendType backend, + size_t q_rank, + uint32_t n_q, uint32_t n_kv, + int64_t is_causal, + float softcap, + bool has_attn_mask, + bool has_qk_output) { + return IsGpuBackend(backend) && + q_rank == 3 && // QNN GQA uses BSH; 4D BNSH → decomposition + n_q > n_kv && n_q % n_kv == 0 && // true GQA per https://arxiv.org/pdf/2305.13245 + is_causal == 1 && // QNN GQA is always causal — no is_causal param + softcap == 0.0f && // no softcap param in QNN GQA + !has_attn_mask && // no additive mask input in QNN GQA + !has_qk_output; // no per-stage debug output in QNN GQA +} + +// Synthesize seqlens_k and total_sequence_length that QNN GQA requires but +// ONNX Attention omits. Both values are derived from static input shapes +// already validated in IsOpSupported. +// +// seqlens_k INT32 [B] = S_past + S_k − 1 (spec: total_seq_len − 1) +// total_seq_len INT32 0D = S_past + S_k +static Ort::Status AddNativeGQASyntheticInputs(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + uint32_t B, uint32_t S_past, uint32_t S_k, + std::string& seqlens_k_name, + std::string& total_seq_len_name) { + const int32_t seqlens_val = static_cast(S_past + S_k) - 1; + const int32_t total_val = static_cast(S_past + S_k); + + seqlens_k_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_seqlens_k"); + { + std::vector bytes(static_cast(B) * sizeof(int32_t)); + int32_t* p = reinterpret_cast(bytes.data()); + for (uint32_t b = 0; b < B; ++b) p[b] = seqlens_val; + QnnTensorWrapper t(seqlens_k_name, QNN_TENSOR_TYPE_STATIC, QNN_DATATYPE_INT_32, + QnnQuantParamsWrapper{}, {B}, std::move(bytes)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(t)), + ("Failed to add seqlens_k: " + seqlens_k_name).c_str()); + } + + total_seq_len_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_total_seq_len"); + { + std::vector bytes(sizeof(int32_t)); + *reinterpret_cast(bytes.data()) = total_val; + // 0D shape (empty dims vector) — QNN requires a scalar, same override used + // by GroupQueryAttentionOpBuilder for com.microsoft::GroupQueryAttention. + QnnTensorWrapper t(total_seq_len_name, QNN_TENSOR_TYPE_STATIC, QNN_DATATYPE_INT_32, + QnnQuantParamsWrapper{}, {}, std::move(bytes)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(t)), + ("Failed to add total_seq_len: " + total_seq_len_name).c_str()); + } + return Ort::Status(); +} + +// Build input_names in the 10-slot QNN GQA order. +// Q/K/V (3D BSH) and past_key/past_value (4D BNSH) are passed straight +// through — no reshape needed. Slots 7-9 (rotary, position) are null. +static Ort::Status ProcessInputsNativeGQA(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger, + std::vector& input_names) { + const auto& onnx_inputs = node_unit.Inputs(); + const bool has_past_key = (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()); + + TensorInfo k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[1], k_info)); + const uint32_t B = k_info.shape[0]; + const uint32_t S_k = k_info.shape[1]; + + uint32_t S_past = 0; + if (has_past_key) { + TensorInfo past_k_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[4], past_k_info)); + S_past = past_k_info.shape[2]; // [B, n_kv, S_past, hs] + } + + auto AddNull = [&](const char* suffix) -> Ort::Status { + const std::string name = utils::UniqueNameGenerator().New(node_unit, suffix); + input_names.push_back(name); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(name)), + ("Failed to add null tensor: " + name).c_str()); + return Ort::Status(); + }; + + // [0] query [1] seqlens_k [2] total_seq_len [3] key [4] value + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[0], logger, input_names)); + + std::string seqlens_k_name, total_seq_name; + RETURN_IF_ERROR(AddNativeGQASyntheticInputs(qnn_model_wrapper, node_unit, + B, S_past, S_k, + seqlens_k_name, total_seq_name)); + input_names.push_back(seqlens_k_name); + input_names.push_back(total_seq_name); + + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[1], logger, input_names)); + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[2], logger, input_names)); + + // [5] past_key [6] past_value — 4D BNSH, passed straight through (matches + // QNN GQA cache format). Null-padded when no KV cache is present. + if (has_past_key) { + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); + RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); + } else { + RETURN_IF_ERROR(AddNull("_null_past_key")); + RETURN_IF_ERROR(AddNull("_null_past_value")); + } + + // [7] cos_cache [8] sin_cache [9] position_ids — no rotary in ai.onnx::Attention + RETURN_IF_ERROR(AddNull("_null_cos")); + RETURN_IF_ERROR(AddNull("_null_sin")); + RETURN_IF_ERROR(AddNull("_null_pos")); + + return Ort::Status(); +} + +// Emit a single QNN_OP_GROUP_QUERY_ATTENTION node (GPU native GQA path). +static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + std::vector&& input_names, + bool do_op_validation) { + const auto& onnx_inputs = node_unit.Inputs(); + const auto& onnx_outputs = node_unit.Outputs(); + OrtNodeAttrHelper node_helper(node_unit); + std::vector param_names; + + // NUM_HEADS + const auto opt_q = node_helper.GetInt64("q_num_heads"); + RETURN_IF_NOT(opt_q.has_value(), "q_num_heads required for native GQA path"); + const uint32_t num_heads_u32 = SafeInt(opt_q.value()); + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), + num_heads_u32, + QNN_OP_GROUP_QUERY_ATTENTION_PARAM_NUM_HEADS, param_names)); + + // KV_NUM_HEADS + const auto opt_kv = node_helper.GetInt64("kv_num_heads"); + RETURN_IF_NOT(opt_kv.has_value(), "kv_num_heads required for native GQA path"); + const uint32_t kv_num_heads_u32 = SafeInt(opt_kv.value()); + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), + kv_num_heads_u32, + QNN_OP_GROUP_QUERY_ATTENTION_PARAM_KV_NUM_HEADS, param_names)); + + // DO_ROTARY = 0 (ai.onnx::Attention has no rotary embeddings) + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), + 0u, + QNN_OP_GROUP_QUERY_ATTENTION_PARAM_DO_ROTARY, + param_names)); + + // SCALE — from attribute or 1/sqrt(head_size); Q is 3D [B, S, n*hs] + TensorInfo q_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[0], q_info)); + const uint32_t head_size = q_info.shape[2] / num_heads_u32; + const float scale_default = 1.0f / std::sqrt(static_cast(head_size)); + const float scale = node_helper.Get("scale", scale_default); + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), + scale, + QNN_OP_GROUP_QUERY_ATTENTION_PARAM_SCALE, param_names)); + + // Outputs: Y (mandatory), present_key, present_value (optional, up to slot 2). + // qk_matmul_output (slot 3) is never reached — rejected by ShouldUseNativeGQA. + std::vector output_names; + const size_t n_outs = std::min(onnx_outputs.size(), size_t{3}); + for (size_t i = 0; i < n_outs; ++i) { + if (onnx_outputs[i].Exists()) { + const std::string& name = onnx_outputs[i].name; + output_names.push_back(name); + TensorInfo out_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[i], out_info)); + const bool is_graph_out = qnn_model_wrapper.IsGraphOutput(name); + QnnTensorWrapper wrapper(name, + is_graph_out ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE, + out_info.qnn_data_type, + std::move(out_info.quant_param), + std::move(out_info.shape)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(wrapper)), + ("Failed to add output: " + name).c_str()); + } else { + const std::string null_name = utils::UniqueNameGenerator().New(node_unit, "_null_out"); + output_names.push_back(null_name); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(null_name)), + "Failed to add null output."); + } + } + + const std::string node_name = utils::UniqueNameGenerator().New(node_unit); + RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_GROUP_QUERY_ATTENTION, + std::move(input_names), + std::move(output_names), + std::move(param_names), + do_op_validation), + "Failed to create QNN_OP_GROUP_QUERY_ATTENTION node."); + return Ort::Status(); +} + // --------------------------------------------------------------------------- // Helper: emit an ElementWiseBinary (MUL or ADD or DIV) node. // --------------------------------------------------------------------------- @@ -685,6 +921,14 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const int64_t is_causal = node_helper.Get("is_causal", static_cast(0)); + // ---- GPU native GQA path ---- + if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), q_rank, + n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output)) { + return EmitNativeGQANode(qnn_model_wrapper, node_unit, + std::move(input_names), do_op_validation); + } + + // ---- Decomposition path ---- const std::string sqrt_scale_name = utils::UniqueNameGenerator().New(node_unit, "_sqrt_scale"); { std::vector scale_bytes; diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index a2424db4ca1..56b3aa84dd9 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -554,6 +554,181 @@ TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode3) { #endif // defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__) +// =========================================================================== +// GPU Attention tests +// +// Two categories: +// Native — GPU + true GQA (n_q > n_kv, divisible) + is_causal=1 +// + no softcap/attn_mask/qk_output → QNN_OP_GROUP_QUERY_ATTENTION +// Decompose — any condition that disqualifies native GQA → decomposition +// (same graph as HTP, but runs on Adreno) +// +// Gated to _M_ARM64: Adreno GPU is only available on Snapdragon X hardware. +// =========================================================================== +#if defined(_M_ARM64) + +// --------------------------------------------------------------------------- +// Native GQA path +// --------------------------------------------------------------------------- + +// GQA 3D, head_ratio=2, causal — all native-GQA conditions met. +// Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8). +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// MQA 3D (kv_num_heads=1 — extreme GQA), causal. +// Q [1,8,32] (4 heads × 8), K/V [1,8,8] (1 head × 8). +TEST_F(QnnGPUBackendTests, Attention_GPU_MQA_3D_Native) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(1)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// GQA 3D with KV cache — past_key/value as static initializers. +// Q [1,8,32], K/V [1,8,16], past_key/value [1,2,4,8] (S_past=4). +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native_KVCache) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCaseKV( + TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // K + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // V + TestInputDef({1, 2, 4, 8}, true, -1.0f, 1.0f), // past_key (initializer) + TestInputDef({1, 2, 4, 8}, true, -1.0f, 1.0f), // past_value (initializer) + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); +} + +// --------------------------------------------------------------------------- +// Decomposition path on GPU +// Each test disqualifies exactly one native-GQA condition. +// --------------------------------------------------------------------------- + +// MHA 3D (n_q == n_kv) — not GQA → decomposition. +TEST_F(QnnGPUBackendTests, Attention_GPU_MHA_3D_Decompose) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// GQA 3D non-causal (is_causal=0) — QNN GQA is always causal → decomposition. +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_NonCausal) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// GQA 3D with softcap — no softcap param in QNN GQA → decomposition. +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(1)), + test::MakeAttribute("softcap", 5.0f)}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// GQA 4D BNSH inputs — QNN GQA uses BSH; 4D falls to decomposition. +// Q [1,4,8,8] (n_q=4), K/V [1,2,8,8] (n_kv=2). +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Decompose) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 8}, false, -1.0f, 1.0f), + TestInputDef({1, 2, 8, 8}, false, -1.0f, 1.0f), + TestInputDef({1, 2, 8, 8}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// GQA 3D with explicit attn_mask — no additive mask in QNN GQA → decomposition. +// mask shape [8,8] broadcast to [1,4,8,8]. +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_AttnMask) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCaseWithMask( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, + TestInputDef({8, 8}, false, -0.5f, 0.0f), // additive float mask + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(2)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +#endif // defined(_M_ARM64) + } // namespace test } // namespace onnxruntime From 9a04ab4bcf0ef092bd224463aa756506f21de78b Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Wed, 29 Jul 2026 18:21:55 -0700 Subject: [PATCH 10/20] Act on review comments --- .../QNN-ExecutionProvider.md | 2 +- .../builder/opbuilder/attention_op_builder.cc | 264 ++++++++++++++---- .../test/providers/qnn/attention_test.cc | 101 +++---- 3 files changed, 237 insertions(+), 130 deletions(-) diff --git a/docs/execution_providers/QNN-ExecutionProvider.md b/docs/execution_providers/QNN-ExecutionProvider.md index 0089ff0140a..e2edfdca082 100644 --- a/docs/execution_providers/QNN-ExecutionProvider.md +++ b/docs/execution_providers/QNN-ExecutionProvider.md @@ -415,7 +415,7 @@ ort.unregister_execution_provider_library(ep_registration_name) |ai.onnx:ArgMin|| |ai.onnx:Asin|| |ai.onnx:Atan|| -|ai.onnx:Attention|Opsets 23–24. 3D \[B,S,n·hs\] and 4D \[B,n,S,hs\] inputs. MHA and GQA/MQA (q\_num\_heads ≥ kv\_num\_heads, divisible). is\_causal, attn\_mask, softcap, KV cache (past/present key·value), qk\_matmul\_output modes 0–3. Static shapes only.| +|ai.onnx:Attention|Opsets 23–24. 3D \[B,S,n·hs\] and 4D \[B,n,S,hs\] inputs. HTP: full decomposition for MHA, GQA/MQA, is\_causal, attn\_mask, softcap, KV cache, qk\_matmul\_output modes 0–3. GPU: native QNN\_OP\_GROUP\_QUERY\_ATTENTION when is\_causal=1 and no softcap/attn\_mask/qk\_output (covers MHA and GQA/MQA); decomposition otherwise. Static shapes only.| |ai.onnx:AveragePool|| |ai.onnx:BatchNormalization|fp16 supported since 1.18.0| |ai.onnx:Cast|| diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 0380ff02018..922e9a33c2d 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -48,6 +48,15 @@ class AttentionOpBuilder : public BaseOpBuilder { std::vector&& input_names, const Ort::Logger& logger, bool do_op_validation) const override ORT_MUST_USE_RESULT; + + private: + // GPU native GQA: ProcessInputsNativeGQA is a static member so it can call + // the protected ProcessInput() inherited from BaseOpBuilder. + static Ort::Status ProcessInputsNativeGQA(const AttentionOpBuilder& self, + QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger, + std::vector& input_names); }; // --------------------------------------------------------------------------- @@ -172,6 +181,12 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper return Ort::Status(); } +// Forward declaration — defined in the GPU native GQA block below. +static bool ShouldUseNativeGQA(QnnBackendType backend, + uint32_t n_q, uint32_t n_kv, + int64_t is_causal, float softcap, + bool has_attn_mask, bool has_qk_output); + // --------------------------------------------------------------------------- // ProcessInputs — register Q, K, V, attn_mask, past_key, past_value, // nonpad_kv_seqlen @@ -205,9 +220,10 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper const bool has_attn_mask = (onnx_inputs.size() > 3 && onnx_inputs[3].Exists()); const bool has_qk_output = (node_unit.Outputs().size() > 3 && node_unit.Outputs()[3].Exists()); - if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), q_rank, - n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output)) { - return ProcessInputsNativeGQA(qnn_model_wrapper, node_unit, logger, input_names); + if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), + n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output, + (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()))) { + return ProcessInputsNativeGQA(*this, qnn_model_wrapper, node_unit, logger, input_names); } } @@ -239,28 +255,31 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper // --------------------------------------------------------------------------- // GPU native GQA path // -// When backend=GPU, input layout is 3D BSH, the node is true GQA -// (n_q > n_kv, divisible), is_causal=1, and no features QNN GQA cannot -// express (softcap/attn_mask/qk_output), a single -// QNN_OP_GROUP_QUERY_ATTENTION node is emitted instead of the decomposed -// graph. All other cases (MHA, non-causal, softcap, HTP, 4D inputs) fall -// through to decomposition unchanged. +// When backend=GPU, kv_num_heads divides num_heads, is_causal=1, and no +// features QNN GQA cannot express (softcap/attn_mask/qk_output), a single +// QNN_OP_GROUP_QUERY_ATTENTION node is emitted. This covers both MHA +// (n_q == n_kv) and GQA/MQA (n_q > n_kv). 4D BNSH inputs are handled by +// inserting Transpose+Reshape before and after the native op. +// All other cases (non-causal, softcap, HTP) fall to decomposition. // --------------------------------------------------------------------------- static bool ShouldUseNativeGQA(QnnBackendType backend, - size_t q_rank, uint32_t n_q, uint32_t n_kv, int64_t is_causal, float softcap, bool has_attn_mask, - bool has_qk_output) { + bool has_qk_output, + bool has_past_key) { + // GPU native GQA requires KV cache (past_key/past_value must be present as + // APP_WRITE tensors). Without a live cache buffer the GPU kernel has no valid + // memory to read or write, causing a runtime access violation. return IsGpuBackend(backend) && - q_rank == 3 && // QNN GQA uses BSH; 4D BNSH → decomposition - n_q > n_kv && n_q % n_kv == 0 && // true GQA per https://arxiv.org/pdf/2305.13245 - is_causal == 1 && // QNN GQA is always causal — no is_causal param - softcap == 0.0f && // no softcap param in QNN GQA - !has_attn_mask && // no additive mask input in QNN GQA - !has_qk_output; // no per-stage debug output in QNN GQA + n_q % n_kv == 0 && // covers MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv) + is_causal == 1 && // QNN GQA is always causal — no is_causal param + softcap == 0.0f && // no softcap param in QNN GQA + !has_attn_mask && // no additive mask input in QNN GQA + !has_qk_output && // no per-stage debug output in QNN GQA + has_past_key; // KV cache required: past_key must be APP_WRITE (dynamic) } // Synthesize seqlens_k and total_sequence_length that QNN GQA requires but @@ -305,10 +324,13 @@ static Ort::Status AddNativeGQASyntheticInputs(QnnModelWrapper& qnn_model_wrappe // Build input_names in the 10-slot QNN GQA order. // Q/K/V (3D BSH) and past_key/past_value (4D BNSH) are passed straight // through — no reshape needed. Slots 7-9 (rotary, position) are null. -static Ort::Status ProcessInputsNativeGQA(QnnModelWrapper& qnn_model_wrapper, - const OrtNodeUnit& node_unit, - const Ort::Logger& logger, - std::vector& input_names) { +// Static member (not free function) so it can call the protected +// ProcessInput() inherited by AttentionOpBuilder from BaseOpBuilder. +Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder& self, + QnnModelWrapper& qnn_model_wrapper, + const OrtNodeUnit& node_unit, + const Ort::Logger& logger, + std::vector& input_names) { const auto& onnx_inputs = node_unit.Inputs(); const bool has_past_key = (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()); @@ -333,7 +355,7 @@ static Ort::Status ProcessInputsNativeGQA(QnnModelWrapper& qnn_model_wrapper, }; // [0] query [1] seqlens_k [2] total_seq_len [3] key [4] value - RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[0], logger, input_names)); + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[0], logger, input_names)); std::string seqlens_k_name, total_seq_name; RETURN_IF_ERROR(AddNativeGQASyntheticInputs(qnn_model_wrapper, node_unit, @@ -342,18 +364,14 @@ static Ort::Status ProcessInputsNativeGQA(QnnModelWrapper& qnn_model_wrapper, input_names.push_back(seqlens_k_name); input_names.push_back(total_seq_name); - RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[1], logger, input_names)); - RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[2], logger, input_names)); + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[1], logger, input_names)); + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[2], logger, input_names)); - // [5] past_key [6] past_value — 4D BNSH, passed straight through (matches - // QNN GQA cache format). Null-padded when no KV cache is present. - if (has_past_key) { - RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); - RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); - } else { - RETURN_IF_ERROR(AddNull("_null_past_key")); - RETURN_IF_ERROR(AddNull("_null_past_value")); - } + // [5] past_key [6] past_value — always present (has_past_key is required by + // ShouldUseNativeGQA). Must be APP_WRITE (dynamic) tensors — the GPU kernel + // requires live cache buffers, not STATIC initializers. + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); // [7] cos_cache [8] sin_cache [9] position_ids — no rotary in ai.onnx::Attention RETURN_IF_ERROR(AddNull("_null_cos")); @@ -363,7 +381,14 @@ static Ort::Status ProcessInputsNativeGQA(QnnModelWrapper& qnn_model_wrapper, return Ort::Status(); } -// Emit a single QNN_OP_GROUP_QUERY_ATTENTION node (GPU native GQA path). +// Emit a single QNN_OP_GROUP_QUERY_ATTENTION node (GPU native path). +// For 4D BNSH inputs: inserts Transpose(0,2,1,3)+Reshape before Q/K/V to +// produce the 3D BSH layout QNN GQA expects, and Reshape+Transpose after +// the Y output to restore the 4D BNSH shape expected by the ONNX graph. +// past_key/past_value are always 4D BNSH and are passed straight through. +// Output shapes are derived from input shapes — not from GetTensorInfo on +// outputs — because the test models do not run ONNX shape inference on +// ai.onnx::Attention and the output value_info has no shape. static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, std::vector&& input_names, @@ -373,7 +398,35 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, OrtNodeAttrHelper node_helper(node_unit); std::vector param_names; - // NUM_HEADS + // Read input shapes to derive all dimensions (no reliance on output shape inference). + TensorInfo q_info{}, v_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[0], q_info)); + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[2], v_info)); + const bool is_4d = (q_info.shape.size() == 4); + const Qnn_DataType_t dtype = q_info.qnn_data_type; + + // Derive tensor dimensions. + uint32_t B = 0, S_q = 0, S_k = 0, S_past = 0; + if (is_4d) { + // Q [B, n_q, S_q, hs] V [B, n_kv, S_k, v_hs] + B = q_info.shape[0]; + S_q = q_info.shape[2]; + S_k = v_info.shape[2]; + } else { + // Q [B, S_q, n_q*hs] V [B, S_k, n_kv*v_hs] + B = q_info.shape[0]; + S_q = q_info.shape[1]; + S_k = v_info.shape[1]; + } + const bool has_past_key = (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()); + if (has_past_key) { + TensorInfo pk_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[4], pk_info)); + S_past = pk_info.shape[2]; // past_key always [B, n_kv, S_past, hs] + } + const uint32_t S_total = S_past + S_k; + + // ---- Params ---- const auto opt_q = node_helper.GetInt64("q_num_heads"); RETURN_IF_NOT(opt_q.has_value(), "q_num_heads required for native GQA path"); const uint32_t num_heads_u32 = SafeInt(opt_q.value()); @@ -381,7 +434,6 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, num_heads_u32, QNN_OP_GROUP_QUERY_ATTENTION_PARAM_NUM_HEADS, param_names)); - // KV_NUM_HEADS const auto opt_kv = node_helper.GetInt64("kv_num_heads"); RETURN_IF_NOT(opt_kv.has_value(), "kv_num_heads required for native GQA path"); const uint32_t kv_num_heads_u32 = SafeInt(opt_kv.value()); @@ -389,48 +441,118 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, kv_num_heads_u32, QNN_OP_GROUP_QUERY_ATTENTION_PARAM_KV_NUM_HEADS, param_names)); - // DO_ROTARY = 0 (ai.onnx::Attention has no rotary embeddings) RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), 0u, QNN_OP_GROUP_QUERY_ATTENTION_PARAM_DO_ROTARY, param_names)); - // SCALE — from attribute or 1/sqrt(head_size); Q is 3D [B, S, n*hs] - TensorInfo q_info{}; - RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[0], q_info)); - const uint32_t head_size = q_info.shape[2] / num_heads_u32; + const uint32_t head_size = is_4d ? q_info.shape[3] + : (q_info.shape[2] / num_heads_u32); + const uint32_t v_hs = is_4d ? v_info.shape[3] + : (v_info.shape[2] / kv_num_heads_u32); const float scale_default = 1.0f / std::sqrt(static_cast(head_size)); const float scale = node_helper.Get("scale", scale_default); RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), scale, QNN_OP_GROUP_QUERY_ATTENTION_PARAM_SCALE, param_names)); - // Outputs: Y (mandatory), present_key, present_value (optional, up to slot 2). - // qk_matmul_output (slot 3) is never reached — rejected by ShouldUseNativeGQA. + // ---- 4D BNSH → 3D BSH transforms for Q (slot 0), K (slot 3), V (slot 4) ---- + if (is_4d) { + constexpr size_t kSlots[3] = {0, 3, 4}; + constexpr size_t kIdx[3] = {0, 1, 2}; + for (int i = 0; i < 3; ++i) { + TensorInfo info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[kIdx[i]], info)); + const uint32_t Bi = info.shape[0], ni = info.shape[1]; + const uint32_t Si = info.shape[2], hsi = info.shape[3]; + const std::string tr_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_tr"); + const std::vector tr_shape = {Bi, Si, ni, hsi}; + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + input_names[kSlots[i]], tr_name, + info.shape, {0u, 2u, 1u, 3u}, + tr_shape, + dtype, info.quant_param, + do_op_validation, false)); + const std::string bsh_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_bsh"); + const std::vector bsh_shape = {Bi, Si, ni * hsi}; + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(tr_name, bsh_name, + tr_shape, bsh_shape, + dtype, info.quant_param, + do_op_validation, false)); + input_names[kSlots[i]] = bsh_name; + } + } + + // ---- Outputs ---- + // All shapes are computed from input shapes (not GetTensorInfo on outputs). std::vector output_names; - const size_t n_outs = std::min(onnx_outputs.size(), size_t{3}); - for (size_t i = 0; i < n_outs; ++i) { - if (onnx_outputs[i].Exists()) { + + // Y: QNN GQA always produces 3D BSH [B, S_q, n_q*v_hs]. + // For 3D input: this is the final ONNX output directly. + // For 4D input: use an intermediate; reshape+transpose back to 4D after the node. + std::string gqa_y_name; + if (onnx_outputs[0].Exists()) { + const std::vector y_bsh = {B, S_q, num_heads_u32 * v_hs}; + if (is_4d) { + gqa_y_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_y_bsh"); + QnnTensorWrapper y3d(gqa_y_name, QNN_TENSOR_TYPE_NATIVE, dtype, + QnnQuantParamsWrapper{}, std::vector(y_bsh)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(y3d)), + "Failed to add GQA Y intermediate tensor."); + } else { + gqa_y_name = onnx_outputs[0].name; + const bool is_go = qnn_model_wrapper.IsGraphOutput(gqa_y_name); + QnnTensorWrapper yw(gqa_y_name, + is_go ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE, + dtype, QnnQuantParamsWrapper{}, std::vector(y_bsh)); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(yw)), + "Failed to add Y output tensor."); + } + } else { + gqa_y_name = utils::UniqueNameGenerator().New(node_unit, "_null_out"); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(gqa_y_name)), + "Failed to add null Y output."); + } + output_names.push_back(gqa_y_name); + + // present_key [B, n_kv, S_total, hs] and present_value [B, n_kv, S_total, v_hs] + // QNN GPU GQA requires real (non-null) tensors for all 3 output slots. + // When the ONNX model does not declare present_key/present_value, allocate + // private NATIVE scratch tensors so QNN can write to them — the caller + // never reads these, but the op validator requires them to exist. + const std::vector pk_shape = {B, kv_num_heads_u32, S_total, head_size}; + const std::vector pv_shape = {B, kv_num_heads_u32, S_total, v_hs}; + const std::vector* cache_shapes[2] = {&pk_shape, &pv_shape}; + for (size_t i = 1; i <= 2; ++i) { + const bool declared = (onnx_outputs.size() > i && onnx_outputs[i].Exists()); + if (declared) { const std::string& name = onnx_outputs[i].name; output_names.push_back(name); - TensorInfo out_info{}; - RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[i], out_info)); - const bool is_graph_out = qnn_model_wrapper.IsGraphOutput(name); - QnnTensorWrapper wrapper(name, - is_graph_out ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE, - out_info.qnn_data_type, - std::move(out_info.quant_param), - std::move(out_info.shape)); - RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(wrapper)), + const bool is_go = qnn_model_wrapper.IsGraphOutput(name); + QnnTensorWrapper cw(name, + is_go ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE, + dtype, QnnQuantParamsWrapper{}, + std::vector(*cache_shapes[i - 1])); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cw)), ("Failed to add output: " + name).c_str()); } else { - const std::string null_name = utils::UniqueNameGenerator().New(node_unit, "_null_out"); - output_names.push_back(null_name); - RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(null_name)), - "Failed to add null output."); + // Private scratch tensor — NATIVE (not APP_READ), never exposed to caller. + const std::string priv_name = utils::UniqueNameGenerator().New(node_unit, "_priv_cache"); + output_names.push_back(priv_name); + QnnTensorWrapper pw(priv_name, QNN_TENSOR_TYPE_NATIVE, dtype, QnnQuantParamsWrapper{}, + std::vector(*cache_shapes[i - 1])); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(pw)), + "Failed to add private cache scratch tensor."); } } + // ---- Emit QNN_OP_GROUP_QUERY_ATTENTION ---- + // Validation is intentionally left on (do_op_validation passed through). + // If past_key/past_value are STATIC initializers the GPU validator returns + // error 3110 and IsOpSupported fails — the node falls to CPU EP rather than + // crashing at runtime. APP_WRITE (dynamic) past_key passes validation and + // takes the native path. seqlens_k/total_seq_len being STATIC is accepted + // by the GPU validator per gpu_v2.json (confirmed by PR #566 tests). const std::string node_name = utils::UniqueNameGenerator().New(node_unit); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, QNN_OP_PACKAGE_NAME_QTI_AISW, @@ -440,6 +562,25 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, std::move(param_names), do_op_validation), "Failed to create QNN_OP_GROUP_QUERY_ATTENTION node."); + + // ---- 3D BSH → 4D BNSH transform for Y (4D input case only) ---- + if (is_4d && onnx_outputs[0].Exists()) { + const std::vector y_bsh = {B, S_q, num_heads_u32 * v_hs}; + const std::vector y_rs = {B, S_q, num_heads_u32, v_hs}; + const std::vector y_bnsh = {B, num_heads_u32, S_q, v_hs}; + const std::string y_rs_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_y_rs"); + RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(gqa_y_name, y_rs_name, + y_bsh, y_rs, + dtype, QnnQuantParamsWrapper{}, + do_op_validation, false)); + const bool is_go = qnn_model_wrapper.IsGraphOutput(onnx_outputs[0].name); + RETURN_IF_ERROR(qnn_model_wrapper.AddTransposeNode(node_unit.Index(), + y_rs_name, onnx_outputs[0].name, + y_rs, {0u, 2u, 1u, 3u}, y_bnsh, + dtype, QnnQuantParamsWrapper{}, + do_op_validation, is_go)); + } + return Ort::Status(); } @@ -922,8 +1063,9 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const int64_t is_causal = node_helper.Get("is_causal", static_cast(0)); // ---- GPU native GQA path ---- - if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), q_rank, - n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output)) { + if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), + n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output, + has_past_key)) { return EmitNativeGQANode(qnn_model_wrapper, node_unit, std::move(input_names), do_op_validation); } diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 56b3aa84dd9..e7c4957322d 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -568,49 +568,13 @@ TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode3) { #if defined(_M_ARM64) // --------------------------------------------------------------------------- -// Native GQA path +// Native GQA path — requires KV cache with APP_WRITE (dynamic) past_key/past_value. +// GPU QNN GQA kernel needs live cache buffers; STATIC initializers are rejected. // --------------------------------------------------------------------------- -// GQA 3D, head_ratio=2, causal — all native-GQA conditions met. -// Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8). -TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native) { - ProviderOptions opts; - opts["backend_type"] = "gpu"; - opts["offload_graph_io_quantization"] = "0"; - - RunQnnModelTest( - BuildAttentionTestCase( - {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), - TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), - TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), - test::MakeAttribute("kv_num_heads", static_cast(2)), - test::MakeAttribute("is_causal", static_cast(1))}), - opts, /*opset_version=*/24, - EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); -} - -// MQA 3D (kv_num_heads=1 — extreme GQA), causal. -// Q [1,8,32] (4 heads × 8), K/V [1,8,8] (1 head × 8). -TEST_F(QnnGPUBackendTests, Attention_GPU_MQA_3D_Native) { - ProviderOptions opts; - opts["backend_type"] = "gpu"; - opts["offload_graph_io_quantization"] = "0"; - - RunQnnModelTest( - BuildAttentionTestCase( - {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), - TestInputDef({1, 8, 8}, false, -1.0f, 1.0f), - TestInputDef({1, 8, 8}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), - test::MakeAttribute("kv_num_heads", static_cast(1)), - test::MakeAttribute("is_causal", static_cast(1))}), - opts, /*opset_version=*/24, - EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); -} - -// GQA 3D with KV cache — past_key/value as static initializers. -// Q [1,8,32], K/V [1,8,16], past_key/value [1,2,4,8] (S_past=4). +// GQA 3D with KV cache, causal — native QNN_OP_GROUP_QUERY_ATTENTION. +// past_key/past_value are dynamic (is_initializer=false) → APP_WRITE as required by GPU. +// Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8), S_past=4. TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native_KVCache) { ProviderOptions opts; opts["backend_type"] = "gpu"; @@ -618,11 +582,11 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native_KVCache) { RunQnnModelTest( BuildAttentionTestCaseKV( - TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q - TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // K - TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // V - TestInputDef({1, 2, 4, 8}, true, -1.0f, 1.0f), // past_key (initializer) - TestInputDef({1, 2, 4, 8}, true, -1.0f, 1.0f), // past_value (initializer) + TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), // Q + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // K + TestInputDef({1, 8, 16}, false, -1.0f, 1.0f), // V + TestInputDef({1, 2, 4, 8}, false, -1.0f, 1.0f), // past_key (dynamic APP_WRITE) + TestInputDef({1, 2, 4, 8}, false, -1.0f, 1.0f), // past_value (dynamic APP_WRITE) {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(2)), test::MakeAttribute("is_causal", static_cast(1))}), @@ -632,29 +596,12 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native_KVCache) { // --------------------------------------------------------------------------- // Decomposition path on GPU -// Each test disqualifies exactly one native-GQA condition. +// Tests without KV cache fall here because ShouldUseNativeGQA requires +// has_past_key=true. Each other test disqualifies one remaining condition. // --------------------------------------------------------------------------- -// MHA 3D (n_q == n_kv) — not GQA → decomposition. -TEST_F(QnnGPUBackendTests, Attention_GPU_MHA_3D_Decompose) { - ProviderOptions opts; - opts["backend_type"] = "gpu"; - opts["offload_graph_io_quantization"] = "0"; - - RunQnnModelTest( - BuildAttentionTestCase( - {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), - TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), - TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("q_num_heads", static_cast(4)), - test::MakeAttribute("kv_num_heads", static_cast(4)), - test::MakeAttribute("is_causal", static_cast(1))}), - opts, /*opset_version=*/24, - EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); -} - -// GQA 3D non-causal (is_causal=0) — QNN GQA is always causal → decomposition. -TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_NonCausal) { +// GQA 3D causal, no KV cache — no past_key → decomposition. +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_NoKVCache) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -690,7 +637,7 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); } -// GQA 4D BNSH inputs — QNN GQA uses BSH; 4D falls to decomposition. +// GQA 4D BNSH inputs, no KV cache — no past_key → decomposition. // Q [1,4,8,8] (n_q=4), K/V [1,2,8,8] (n_kv=2). TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Decompose) { ProviderOptions opts; @@ -727,6 +674,24 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_AttnMask) { EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); } +// MHA 3D non-causal (is_causal=0) — causal required for native → decomposition. +TEST_F(QnnGPUBackendTests, Attention_GPU_MHA_3D_Decompose_NonCausal) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(0))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + #endif // defined(_M_ARM64) } // namespace test From b7553f35abb13d18adf7831ccdfeceb642b04a5d Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Wed, 29 Jul 2026 18:51:04 -0700 Subject: [PATCH 11/20] CI fixes --- .../providers/qnn/builder/opbuilder/attention_op_builder.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 922e9a33c2d..d567f122bc1 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -185,7 +185,8 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper static bool ShouldUseNativeGQA(QnnBackendType backend, uint32_t n_q, uint32_t n_kv, int64_t is_causal, float softcap, - bool has_attn_mask, bool has_qk_output); + bool has_attn_mask, bool has_qk_output, + bool has_past_key); // --------------------------------------------------------------------------- // ProcessInputs — register Q, K, V, attn_mask, past_key, past_value, From 3c15d602c3ed0d730d39601a7f46a2ad7d6ffc46 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Wed, 29 Jul 2026 21:40:31 -0700 Subject: [PATCH 12/20] [QNN EP] Fix build failures --- .../providers/qnn/builder/opbuilder/attention_op_builder.cc | 6 +++++- onnxruntime/test/providers/qnn/attention_test.cc | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index d567f122bc1..2f652afe9ff 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -350,7 +350,11 @@ Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder& auto AddNull = [&](const char* suffix) -> Ort::Status { const std::string name = utils::UniqueNameGenerator().New(node_unit, suffix); input_names.push_back(name); - RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(name)), + // Use QNN_DATATYPE_UNDEFINED — matches GroupQueryAttentionOpBuilder's null tensor + // pattern. QNN_DATATYPE_FLOAT_32 (from MakeNull) causes validator error 3110. + QnnTensorWrapper null_wrapper(name, QNN_TENSOR_TYPE_NULL, QNN_DATATYPE_UNDEFINED, + QnnQuantParamsWrapper(), std::vector{0}); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(null_wrapper)), ("Failed to add null tensor: " + name).c_str()); return Ort::Status(); }; diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index e7c4957322d..80605680c62 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -634,7 +634,7 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { test::MakeAttribute("is_causal", static_cast(1)), test::MakeAttribute("softcap", 5.0f)}), opts, /*opset_version=*/24, - EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-2f)}); } // GQA 4D BNSH inputs, no KV cache — no past_key → decomposition. From 8411dd9989193fffea5d60f5f38f50e872f353da Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 30 Jul 2026 14:09:34 -0700 Subject: [PATCH 13/20] Act on review comments --- .../builder/opbuilder/attention_op_builder.cc | 153 +++++++++--------- .../test/providers/qnn/attention_test.cc | 87 ++++++++-- 2 files changed, 151 insertions(+), 89 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 2f652afe9ff..8f2d1c179f6 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -22,10 +22,8 @@ namespace qnn { // 3D inputs [B, S, n*hs] (BSH layout, reshape + transpose to BNSH) // GQA/MQA, KV cache (past/present), softcap, qk_matmul_output // -// The SDK version guard matches GroupQueryAttentionOpBuilder so that this class -// is compiled away when building against an SDK without the required op-set -// version metadata. -#if !(QNN_OPSET_VERSION_MAJOR < 2 || (QNN_OPSET_VERSION_MAJOR == 2 && QNN_OPSET_VERSION_MINOR <= 11)) +// The decomposition path is always available; the GPU native path (which emits +// QNN_OP_GROUP_QUERY_ATTENTION) is gated to SDK >= 2.12 (QNN opset 2.12). class AttentionOpBuilder : public BaseOpBuilder { public: @@ -50,8 +48,6 @@ class AttentionOpBuilder : public BaseOpBuilder { bool do_op_validation) const override ORT_MUST_USE_RESULT; private: - // GPU native GQA: ProcessInputsNativeGQA is a static member so it can call - // the protected ProcessInput() inherited from BaseOpBuilder. static Ort::Status ProcessInputsNativeGQA(const AttentionOpBuilder& self, QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, @@ -116,25 +112,11 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper RETURN_IF(q_rank != 3 && q_rank != 4, "Attention: Q input must be rank 3 ([B,S_q,n_q*hs]) or rank 4 ([B,n_q,S_q,hs])"); - // ---- Reject dynamic (0-dim) shapes on Q, K, V ---- - for (uint32_t d : q_info.shape) { - RETURN_IF(d == 0, - "Attention: Q input contains a dynamic (0) dimension; only static shapes are supported"); - } - TensorInfo k_info{}; RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[1], k_info)); - for (uint32_t d : k_info.shape) { - RETURN_IF(d == 0, - "Attention: K input contains a dynamic (0) dimension; only static shapes are supported"); - } TensorInfo v_info{}; RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[2], v_info)); - for (uint32_t d : v_info.shape) { - RETURN_IF(d == 0, - "Attention: V input contains a dynamic (0) dimension; only static shapes are supported"); - } // ---- For 3D inputs, q_num_heads and kv_num_heads attrs are required ---- if (q_rank == 3) { @@ -167,10 +149,6 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[4], past_k_info)); RETURN_IF(past_k_info.shape.size() != 4, "Attention: past_key must be rank 4 ([B,n,S_past,hs])"); - for (uint32_t d : past_k_info.shape) { - RETURN_IF(d == 0, - "Attention: past_key contains a dynamic (0) dimension; only static shapes are supported"); - } } // ---- Full validation: build decomposed nodes with do_op_validation=true ---- @@ -253,9 +231,11 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper return Ort::Status(); } +#if !(QNN_OPSET_VERSION_MAJOR < 2 || (QNN_OPSET_VERSION_MAJOR == 2 && QNN_OPSET_VERSION_MINOR <= 11)) // --------------------------------------------------------------------------- -// GPU native GQA path +// GPU native GQA path (SDK >= 2.12 only) // +// QNN_OP_GROUP_QUERY_ATTENTION is only available in QNN opset 2.12+. // When backend=GPU, kv_num_heads divides num_heads, is_causal=1, and no // features QNN GQA cannot express (softcap/attn_mask/qk_output), a single // QNN_OP_GROUP_QUERY_ATTENTION node is emitted. This covers both MHA @@ -270,17 +250,13 @@ static bool ShouldUseNativeGQA(QnnBackendType backend, float softcap, bool has_attn_mask, bool has_qk_output, - bool has_past_key) { - // GPU native GQA requires KV cache (past_key/past_value must be present as - // APP_WRITE tensors). Without a live cache buffer the GPU kernel has no valid - // memory to read or write, causing a runtime access violation. + bool /*has_past_key*/) { return IsGpuBackend(backend) && n_q % n_kv == 0 && // covers MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv) is_causal == 1 && // QNN GQA is always causal — no is_causal param softcap == 0.0f && // no softcap param in QNN GQA !has_attn_mask && // no additive mask input in QNN GQA - !has_qk_output && // no per-stage debug output in QNN GQA - has_past_key; // KV cache required: past_key must be APP_WRITE (dynamic) + !has_qk_output; // no per-stage debug output in QNN GQA } // Synthesize seqlens_k and total_sequence_length that QNN GQA requires but @@ -312,8 +288,11 @@ static Ort::Status AddNativeGQASyntheticInputs(QnnModelWrapper& qnn_model_wrappe { std::vector bytes(sizeof(int32_t)); *reinterpret_cast(bytes.data()) = total_val; - // 0D shape (empty dims vector) — QNN requires a scalar, same override used - // by GroupQueryAttentionOpBuilder for com.microsoft::GroupQueryAttention. + // 0D shape (empty dims vector) — QNN requires a scalar. + // TODO: GetOnnxShape in qnn_model_wrapper.cc forces scalars to rank 1; + // if total_seq_len ever came from the ONNX graph that shape override pattern + // would be needed here too. Since we synthesize this tensor ourselves we + // avoid that issue, but the root cause should be fixed in GetOnnxShape. QnnTensorWrapper t(total_seq_len_name, QNN_TENSOR_TYPE_STATIC, QNN_DATATYPE_INT_32, QnnQuantParamsWrapper{}, {}, std::move(bytes)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(t)), @@ -350,11 +329,7 @@ Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder& auto AddNull = [&](const char* suffix) -> Ort::Status { const std::string name = utils::UniqueNameGenerator().New(node_unit, suffix); input_names.push_back(name); - // Use QNN_DATATYPE_UNDEFINED — matches GroupQueryAttentionOpBuilder's null tensor - // pattern. QNN_DATATYPE_FLOAT_32 (from MakeNull) causes validator error 3110. - QnnTensorWrapper null_wrapper(name, QNN_TENSOR_TYPE_NULL, QNN_DATATYPE_UNDEFINED, - QnnQuantParamsWrapper(), std::vector{0}); - RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(null_wrapper)), + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(name)), ("Failed to add null tensor: " + name).c_str()); return Ort::Status(); }; @@ -372,11 +347,15 @@ Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder& RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[1], logger, input_names)); RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[2], logger, input_names)); - // [5] past_key [6] past_value — always present (has_past_key is required by - // ShouldUseNativeGQA). Must be APP_WRITE (dynamic) tensors — the GPU kernel - // requires live cache buffers, not STATIC initializers. - RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); - RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); + // [5] past_key [6] past_value — 4D BNSH, passed straight through when present. + // Null-padded when no KV cache (has_past_key=false). + if (has_past_key) { + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); + RETURN_IF_ERROR(self.ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); + } else { + RETURN_IF_ERROR(AddNull("_null_past_key")); + RETURN_IF_ERROR(AddNull("_null_past_value")); + } // [7] cos_cache [8] sin_cache [9] position_ids — no rotary in ai.onnx::Attention RETURN_IF_ERROR(AddNull("_null_cos")); @@ -391,9 +370,6 @@ Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder& // produce the 3D BSH layout QNN GQA expects, and Reshape+Transpose after // the Y output to restore the 4D BNSH shape expected by the ONNX graph. // past_key/past_value are always 4D BNSH and are passed straight through. -// Output shapes are derived from input shapes — not from GetTensorInfo on -// outputs — because the test models do not run ONNX shape inference on -// ai.onnx::Attention and the output value_info has no shape. static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, std::vector&& input_names, @@ -432,16 +408,27 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, const uint32_t S_total = S_past + S_k; // ---- Params ---- - const auto opt_q = node_helper.GetInt64("q_num_heads"); - RETURN_IF_NOT(opt_q.has_value(), "q_num_heads required for native GQA path"); - const uint32_t num_heads_u32 = SafeInt(opt_q.value()); + // For 4D inputs [B, n_q, S_q, hs] the head counts are implicit in the shape; + // for 3D inputs [B, S, n*hs] they must be provided as attributes. + uint32_t num_heads_u32 = 0; + uint32_t kv_num_heads_u32 = 0; + if (is_4d) { + TensorInfo k_info_4d{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[1], k_info_4d)); + num_heads_u32 = q_info.shape[1]; + kv_num_heads_u32 = k_info_4d.shape[1]; + } else { + const auto opt_q = node_helper.GetInt64("q_num_heads"); + RETURN_IF_NOT(opt_q.has_value(), "q_num_heads attribute required for 3D native GQA path"); + const auto opt_kv = node_helper.GetInt64("kv_num_heads"); + RETURN_IF_NOT(opt_kv.has_value(), "kv_num_heads attribute required for 3D native GQA path"); + num_heads_u32 = SafeInt(opt_q.value()); + kv_num_heads_u32 = SafeInt(opt_kv.value()); + } + RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), num_heads_u32, QNN_OP_GROUP_QUERY_ATTENTION_PARAM_NUM_HEADS, param_names)); - - const auto opt_kv = node_helper.GetInt64("kv_num_heads"); - RETURN_IF_NOT(opt_kv.has_value(), "kv_num_heads required for native GQA path"); - const uint32_t kv_num_heads_u32 = SafeInt(opt_kv.value()); RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, node_unit.Index(), node_unit.Name(), kv_num_heads_u32, QNN_OP_GROUP_QUERY_ATTENTION_PARAM_KV_NUM_HEADS, param_names)); @@ -489,7 +476,6 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, } // ---- Outputs ---- - // All shapes are computed from input shapes (not GetTensorInfo on outputs). std::vector output_names; // Y: QNN GQA always produces 3D BSH [B, S_q, n_q*v_hs]. @@ -497,10 +483,12 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, // For 4D input: use an intermediate; reshape+transpose back to 4D after the node. std::string gqa_y_name; if (onnx_outputs[0].Exists()) { - const std::vector y_bsh = {B, S_q, num_heads_u32 * v_hs}; + TensorInfo y_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[0], y_info)); if (is_4d) { gqa_y_name = utils::UniqueNameGenerator().New(node_unit, "_gqa_y_bsh"); - QnnTensorWrapper y3d(gqa_y_name, QNN_TENSOR_TYPE_NATIVE, dtype, + const std::vector y_bsh = {B, S_q, num_heads_u32 * v_hs}; + QnnTensorWrapper y3d(gqa_y_name, QNN_TENSOR_TYPE_NATIVE, y_info.qnn_data_type, QnnQuantParamsWrapper{}, std::vector(y_bsh)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(y3d)), "Failed to add GQA Y intermediate tensor."); @@ -509,7 +497,8 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, const bool is_go = qnn_model_wrapper.IsGraphOutput(gqa_y_name); QnnTensorWrapper yw(gqa_y_name, is_go ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE, - dtype, QnnQuantParamsWrapper{}, std::vector(y_bsh)); + y_info.qnn_data_type, std::move(y_info.quant_param), + std::move(y_info.shape)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(yw)), "Failed to add Y output tensor."); } @@ -520,11 +509,11 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, } output_names.push_back(gqa_y_name); - // present_key [B, n_kv, S_total, hs] and present_value [B, n_kv, S_total, v_hs] + // present_key and present_value (output slots 1-2). // QNN GPU GQA requires real (non-null) tensors for all 3 output slots. // When the ONNX model does not declare present_key/present_value, allocate // private NATIVE scratch tensors so QNN can write to them — the caller - // never reads these, but the op validator requires them to exist. + // never reads these, but the QNN op def for GQA requires them. const std::vector pk_shape = {B, kv_num_heads_u32, S_total, head_size}; const std::vector pv_shape = {B, kv_num_heads_u32, S_total, v_hs}; const std::vector* cache_shapes[2] = {&pk_shape, &pv_shape}; @@ -533,11 +522,13 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, if (declared) { const std::string& name = onnx_outputs[i].name; output_names.push_back(name); + TensorInfo out_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_outputs[i], out_info)); const bool is_go = qnn_model_wrapper.IsGraphOutput(name); QnnTensorWrapper cw(name, is_go ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE, - dtype, QnnQuantParamsWrapper{}, - std::vector(*cache_shapes[i - 1])); + out_info.qnn_data_type, std::move(out_info.quant_param), + std::move(out_info.shape)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cw)), ("Failed to add output: " + name).c_str()); } else { @@ -552,12 +543,9 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, } // ---- Emit QNN_OP_GROUP_QUERY_ATTENTION ---- - // Validation is intentionally left on (do_op_validation passed through). - // If past_key/past_value are STATIC initializers the GPU validator returns - // error 3110 and IsOpSupported fails — the node falls to CPU EP rather than - // crashing at runtime. APP_WRITE (dynamic) past_key passes validation and - // takes the native path. seqlens_k/total_seq_len being STATIC is accepted - // by the GPU validator per gpu_v2.json (confirmed by PR #566 tests). + // Validation is intentionally left on (do_op_validation passed through) so + // that any unsupported configuration is caught at IsOpSupported time and the + // node falls to CPU EP gracefully, rather than failing at runtime. const std::string node_name = utils::UniqueNameGenerator().New(node_unit); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, QNN_OP_PACKAGE_NAME_QTI_AISW, @@ -589,6 +577,28 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, return Ort::Status(); } +#else // SDK < 2.12: no QNN_OP_GROUP_QUERY_ATTENTION — always use decomposition + +static bool ShouldUseNativeGQA(QnnBackendType, uint32_t, uint32_t, + int64_t, float, bool, bool, bool) { + return false; +} + +Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder&, + QnnModelWrapper&, + const OrtNodeUnit&, + const Ort::Logger&, + std::vector&) { + return Ort::Status(); +} + +static Ort::Status EmitNativeGQANode(QnnModelWrapper&, const OrtNodeUnit&, + std::vector&&, bool) { + return Ort::Status(); +} + +#endif // SDK version guard for QNN_OP_GROUP_QUERY_ATTENTION + // --------------------------------------------------------------------------- // Helper: emit an ElementWiseBinary (MUL or ADD or DIV) node. // --------------------------------------------------------------------------- @@ -1535,14 +1545,5 @@ void CreateAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations op_registrations.AddOpBuilder(op_type, std::make_unique()); } -#else // SDK version guard - -void CreateAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { - ORT_UNUSED_PARAMETER(op_type); - ORT_UNUSED_PARAMETER(op_registrations); -} - -#endif // !(QNN_OPSET_VERSION_MAJOR < 2 || (QNN_OPSET_VERSION_MAJOR == 2 && QNN_OPSET_VERSION_MINOR <= 11)) - } // namespace qnn } // namespace onnxruntime diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 80605680c62..6450d51bc05 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -568,14 +568,14 @@ TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode3) { #if defined(_M_ARM64) // --------------------------------------------------------------------------- -// Native GQA path — requires KV cache with APP_WRITE (dynamic) past_key/past_value. -// GPU QNN GQA kernel needs live cache buffers; STATIC initializers are rejected. +// Native GQA path // --------------------------------------------------------------------------- // GQA 3D with KV cache, causal — native QNN_OP_GROUP_QUERY_ATTENTION. -// past_key/past_value are dynamic (is_initializer=false) → APP_WRITE as required by GPU. -// Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8), S_past=4. -TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native_KVCache) { +// DISABLED: unpacked QKV (separate K/V inputs) is not supported by the GPU backend +// in QAIRT 2.48; expected to be supported in QAIRT 2.50. Re-enable when the SDK +// is uplevelled. +TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Native_KVCache) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -595,13 +595,13 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native_KVCache) { } // --------------------------------------------------------------------------- -// Decomposition path on GPU -// Tests without KV cache fall here because ShouldUseNativeGQA requires -// has_past_key=true. Each other test disqualifies one remaining condition. +// Native GQA path (GPU + causal + no softcap/attn_mask/qk_output) // --------------------------------------------------------------------------- -// GQA 3D causal, no KV cache — no past_key → decomposition. -TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_NoKVCache) { +// GQA 3D, head_ratio=2, causal, no KV cache → native QNN_OP_GROUP_QUERY_ATTENTION. +// Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8). +// DISABLED: unpacked QKV not supported by GPU backend in QAIRT 2.48; re-enable in 2.50. +TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Native) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -613,11 +613,71 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_NoKVCache) { TestInputDef({1, 8, 16}, false, -1.0f, 1.0f)}, {test::MakeAttribute("q_num_heads", static_cast(4)), test::MakeAttribute("kv_num_heads", static_cast(2)), - test::MakeAttribute("is_causal", static_cast(0))}), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// MQA 3D (kv_num_heads=1), causal, no KV cache → native path. +// DISABLED: unpacked QKV not supported by GPU backend in QAIRT 2.48; re-enable in 2.50. +TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_MQA_3D_Native) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 32}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 8}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(1)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// MHA 3D (n_q == n_kv), causal, no KV cache → native path (MHA is supported). +// DISABLED: unpacked QKV not supported by GPU backend in QAIRT 2.48; re-enable in 2.50. +TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_MHA_3D_Native) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f), + TestInputDef({1, 8, 64}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("q_num_heads", static_cast(4)), + test::MakeAttribute("kv_num_heads", static_cast(4)), + test::MakeAttribute("is_causal", static_cast(1))}), + opts, /*opset_version=*/24, + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); +} + +// GQA 4D BNSH, causal, no KV cache → native (EmitNativeGQANode inserts Transpose+Reshape). +// DISABLED: after 4D→3D transforms the QNN node has unpacked K/V inputs, which is not +// supported by the GPU backend in QAIRT 2.48; re-enable in 2.50. +TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_4D_Native) { + ProviderOptions opts; + opts["backend_type"] = "gpu"; + opts["offload_graph_io_quantization"] = "0"; + + RunQnnModelTest( + BuildAttentionTestCase( + {TestInputDef({1, 4, 8, 8}, false, -1.0f, 1.0f), + TestInputDef({1, 2, 8, 8}, false, -1.0f, 1.0f), + TestInputDef({1, 2, 8, 8}, false, -1.0f, 1.0f)}, + {test::MakeAttribute("is_causal", static_cast(1))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); } +// --------------------------------------------------------------------------- +// Decomposition path on GPU — each test disqualifies one native condition. +// --------------------------------------------------------------------------- + // GQA 3D with softcap — no softcap param in QNN GQA → decomposition. TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { ProviderOptions opts; @@ -634,10 +694,11 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { test::MakeAttribute("is_causal", static_cast(1)), test::MakeAttribute("softcap", 5.0f)}), opts, /*opset_version=*/24, - EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-2f)}); + EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); } // GQA 4D BNSH inputs, no KV cache — no past_key → decomposition. +// GQA 4D BNSH inputs, no KV cache → decomposition (is_causal=0 disqualifies native path). // Q [1,4,8,8] (n_q=4), K/V [1,2,8,8] (n_kv=2). TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Decompose) { ProviderOptions opts; @@ -649,7 +710,7 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Decompose) { {TestInputDef({1, 4, 8, 8}, false, -1.0f, 1.0f), TestInputDef({1, 2, 8, 8}, false, -1.0f, 1.0f), TestInputDef({1, 2, 8, 8}, false, -1.0f, 1.0f)}, - {test::MakeAttribute("is_causal", static_cast(1))}), + {test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); } From f522c24a8f1ef5b5e27d0fd96091cd7cbff73e89 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 30 Jul 2026 15:43:04 -0700 Subject: [PATCH 14/20] Fixes --- .../builder/opbuilder/attention_op_builder.cc | 43 ++++++++++--------- .../test/providers/qnn/attention_test.cc | 16 +++---- 2 files changed, 31 insertions(+), 28 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 8f2d1c179f6..d989e5a6d82 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -164,7 +164,7 @@ static bool ShouldUseNativeGQA(QnnBackendType backend, uint32_t n_q, uint32_t n_kv, int64_t is_causal, float softcap, bool has_attn_mask, bool has_qk_output, - bool has_past_key); + bool has_present_key); // --------------------------------------------------------------------------- // ProcessInputs — register Q, K, V, attn_mask, past_key, past_value, @@ -201,7 +201,7 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper node_unit.Outputs()[3].Exists()); if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output, - (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()))) { + (node_unit.Outputs().size() > 1 && node_unit.Outputs()[1].Exists()))) { return ProcessInputsNativeGQA(*this, qnn_model_wrapper, node_unit, logger, input_names); } } @@ -241,7 +241,7 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper // QNN_OP_GROUP_QUERY_ATTENTION node is emitted. This covers both MHA // (n_q == n_kv) and GQA/MQA (n_q > n_kv). 4D BNSH inputs are handled by // inserting Transpose+Reshape before and after the native op. -// All other cases (non-causal, softcap, HTP) fall to decomposition. +// All other cases (non-causal, softcap, no KV cache, HTP) fall to decomposition. // --------------------------------------------------------------------------- static bool ShouldUseNativeGQA(QnnBackendType backend, @@ -250,13 +250,17 @@ static bool ShouldUseNativeGQA(QnnBackendType backend, float softcap, bool has_attn_mask, bool has_qk_output, - bool /*has_past_key*/) { + bool has_present_key) { + // TODO: Remove has_present_key once GPU backend adds support for absent KV + // cache outputs (currently present_key/present_value are required by the GPU + // validator even though the QNN op def marks them as optional). return IsGpuBackend(backend) && - n_q % n_kv == 0 && // covers MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv) - is_causal == 1 && // QNN GQA is always causal — no is_causal param - softcap == 0.0f && // no softcap param in QNN GQA - !has_attn_mask && // no additive mask input in QNN GQA - !has_qk_output; // no per-stage debug output in QNN GQA + n_q % n_kv == 0 && // covers MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv) + is_causal == 1 && // QNN GQA is always causal — no is_causal param + softcap == 0.0f && // no softcap param in QNN GQA + !has_attn_mask && // no additive mask input in QNN GQA + !has_qk_output && // no per-stage debug output in QNN GQA + has_present_key; // GPU validator currently requires KV cache outputs } // Synthesize seqlens_k and total_sequence_length that QNN GQA requires but @@ -514,9 +518,8 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, // When the ONNX model does not declare present_key/present_value, allocate // private NATIVE scratch tensors so QNN can write to them — the caller // never reads these, but the QNN op def for GQA requires them. - const std::vector pk_shape = {B, kv_num_heads_u32, S_total, head_size}; - const std::vector pv_shape = {B, kv_num_heads_u32, S_total, v_hs}; - const std::vector* cache_shapes[2] = {&pk_shape, &pv_shape}; + // present_key and present_value (output slots 1-2). + // ShouldUseNativeGQA requires has_present_key=true so both are always declared. for (size_t i = 1; i <= 2; ++i) { const bool declared = (onnx_outputs.size() > i && onnx_outputs[i].Exists()); if (declared) { @@ -532,13 +535,13 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cw)), ("Failed to add output: " + name).c_str()); } else { - // Private scratch tensor — NATIVE (not APP_READ), never exposed to caller. - const std::string priv_name = utils::UniqueNameGenerator().New(node_unit, "_priv_cache"); - output_names.push_back(priv_name); - QnnTensorWrapper pw(priv_name, QNN_TENSOR_TYPE_NATIVE, dtype, QnnQuantParamsWrapper{}, - std::vector(*cache_shapes[i - 1])); - RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(pw)), - "Failed to add private cache scratch tensor."); + // Absent cache output — use null tensor. ShouldUseNativeGQA ensures + // present_key is always declared, so only present_value can reach here + // if the model omits it. + const std::string null_name = utils::UniqueNameGenerator().New(node_unit, "_null_out"); + output_names.push_back(null_name); + RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(QnnTensorWrapper::MakeNull(null_name)), + "Failed to add null cache output."); } } @@ -1080,7 +1083,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // ---- GPU native GQA path ---- if (ShouldUseNativeGQA(qnn_model_wrapper.GetQnnBackendType(), n_q, n_kv, is_causal, softcap, has_attn_mask, has_qk_output, - has_past_key)) { + (onnx_outputs.size() > 1 && onnx_outputs[1].Exists()))) { return EmitNativeGQANode(qnn_model_wrapper, node_unit, std::move(input_names), do_op_validation); } diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 6450d51bc05..ed46060a0b0 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -600,8 +600,7 @@ TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Native_KVCache) { // GQA 3D, head_ratio=2, causal, no KV cache → native QNN_OP_GROUP_QUERY_ATTENTION. // Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8). -// DISABLED: unpacked QKV not supported by GPU backend in QAIRT 2.48; re-enable in 2.50. -TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Native) { +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -619,8 +618,7 @@ TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Native) { } // MQA 3D (kv_num_heads=1), causal, no KV cache → native path. -// DISABLED: unpacked QKV not supported by GPU backend in QAIRT 2.48; re-enable in 2.50. -TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_MQA_3D_Native) { +TEST_F(QnnGPUBackendTests, Attention_GPU_MQA_3D_Native) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -638,8 +636,7 @@ TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_MQA_3D_Native) { } // MHA 3D (n_q == n_kv), causal, no KV cache → native path (MHA is supported). -// DISABLED: unpacked QKV not supported by GPU backend in QAIRT 2.48; re-enable in 2.50. -TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_MHA_3D_Native) { +TEST_F(QnnGPUBackendTests, Attention_GPU_MHA_3D_Native) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -659,7 +656,7 @@ TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_MHA_3D_Native) { // GQA 4D BNSH, causal, no KV cache → native (EmitNativeGQANode inserts Transpose+Reshape). // DISABLED: after 4D→3D transforms the QNN node has unpacked K/V inputs, which is not // supported by the GPU backend in QAIRT 2.48; re-enable in 2.50. -TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_4D_Native) { +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Native) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; @@ -679,7 +676,10 @@ TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_4D_Native) { // --------------------------------------------------------------------------- // GQA 3D with softcap — no softcap param in QNN GQA → decomposition. -TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { +// DISABLED: accuracy failure on GPU backend +// Max observed delta: ~0.028. +// Tracked as a GPU backend precision issue. +TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Decompose_Softcap) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; From 8c6fc24c3005d0992d4de9b476744b0a6e6353f8 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 30 Jul 2026 19:32:15 -0700 Subject: [PATCH 15/20] [QNN EP] Fix CI build failures --- .../qnn/builder/opbuilder/attention_op_builder.cc | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index d989e5a6d82..0ec11646e7b 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -255,12 +255,12 @@ static bool ShouldUseNativeGQA(QnnBackendType backend, // cache outputs (currently present_key/present_value are required by the GPU // validator even though the QNN op def marks them as optional). return IsGpuBackend(backend) && - n_q % n_kv == 0 && // covers MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv) - is_causal == 1 && // QNN GQA is always causal — no is_causal param - softcap == 0.0f && // no softcap param in QNN GQA - !has_attn_mask && // no additive mask input in QNN GQA - !has_qk_output && // no per-stage debug output in QNN GQA - has_present_key; // GPU validator currently requires KV cache outputs + n_q % n_kv == 0 && // covers MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv) + is_causal == 1 && // QNN GQA is always causal — no is_causal param + softcap == 0.0f && // no softcap param in QNN GQA + !has_attn_mask && // no additive mask input in QNN GQA + !has_qk_output && // no per-stage debug output in QNN GQA + has_present_key; // GPU validator currently requires KV cache outputs } // Synthesize seqlens_k and total_sequence_length that QNN GQA requires but @@ -409,7 +409,6 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[4], pk_info)); S_past = pk_info.shape[2]; // past_key always [B, n_kv, S_past, hs] } - const uint32_t S_total = S_past + S_k; // ---- Params ---- // For 4D inputs [B, n_q, S_q, hs] the head counts are implicit in the shape; From cb80676bd8153bf24a8a7f741983891f1905fab6 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 30 Jul 2026 20:10:04 -0700 Subject: [PATCH 16/20] CI fix --- .../qnn/builder/opbuilder/attention_op_builder.cc | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 0ec11646e7b..cbaf96597b0 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -390,24 +390,14 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, const bool is_4d = (q_info.shape.size() == 4); const Qnn_DataType_t dtype = q_info.qnn_data_type; - // Derive tensor dimensions. - uint32_t B = 0, S_q = 0, S_k = 0, S_past = 0; + // Derive tensor dimensions needed for 4D output transforms. + uint32_t B = 0, S_q = 0; if (is_4d) { - // Q [B, n_q, S_q, hs] V [B, n_kv, S_k, v_hs] B = q_info.shape[0]; S_q = q_info.shape[2]; - S_k = v_info.shape[2]; } else { - // Q [B, S_q, n_q*hs] V [B, S_k, n_kv*v_hs] B = q_info.shape[0]; S_q = q_info.shape[1]; - S_k = v_info.shape[1]; - } - const bool has_past_key = (onnx_inputs.size() > 4 && onnx_inputs[4].Exists()); - if (has_past_key) { - TensorInfo pk_info{}; - RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(onnx_inputs[4], pk_info)); - S_past = pk_info.shape[2]; // past_key always [B, n_kv, S_past, hs] } // ---- Params ---- From 4e2b319972a05be32b17d59146bfd8c61696b036 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Wed, 5 Aug 2026 14:31:55 -0700 Subject: [PATCH 17/20] [QNN EP] Act on review comments --- .../builder/opbuilder/attention_op_builder.cc | 40 +++++++++---------- .../test/providers/qnn/attention_test.cc | 6 +-- 2 files changed, 21 insertions(+), 25 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index cbaf96597b0..bb3af4bf330 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -1302,10 +1302,24 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn qk_captured = scores_cur; } - // ---- Masks applied BEFORE softcap (per ONNX spec) ---- - // Applying mask after softcap would be wrong: softcap saturates large negatives - // to -softcap (not -inf), so the mask must set large negatives first so softcap - // can compress the full range uniformly. + // ---- Softcap BEFORE masks (per ONNX spec) ---- + // From the ONNX spec function body: MatMul → softcap → attn_mask+Add → Softmax. + // Softcap must come first: applying the mask (-1e9) before softcap would clamp + // it to -softcap (e.g. -5.0), making masked positions visible to softmax. + if (softcap != 0.0f) { + const std::string sc_out = utils::UniqueNameGenerator().New(node_unit, "_scores_softcap"); + RETURN_IF_ERROR(AddSoftcapNode(qnn_model_wrapper, node_unit, + scores_cur, sc_out, + qk_shape, dtype, q_quant, + softcap, + do_op_validation)); + scores_cur = sc_out; + } + + // ---- qk_matmul_output mode 1: post-softcap, pre-mask ---- + if (has_qk_output && qk_mode == 1) { + qk_captured = scores_cur; + } // ---- Causal mask (is_causal=1): ADD static lower-triangular mask ---- // With KV cache: offset = S_past so that row i attends to positions <= i+S_past. @@ -1380,23 +1394,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn scores_cur = attn_masked_out; } - // ---- qk_matmul_output mode 1: post-mask, pre-softcap ---- - if (has_qk_output && qk_mode == 1) { - qk_captured = scores_cur; - } - - // ---- Softcap (applied AFTER masks per ONNX spec) ---- - if (softcap != 0.0f) { - const std::string sc_out = utils::UniqueNameGenerator().New(node_unit, "_scores_softcap"); - RETURN_IF_ERROR(AddSoftcapNode(qnn_model_wrapper, node_unit, - scores_cur, sc_out, - qk_shape, dtype, q_quant, - softcap, - do_op_validation)); - scores_cur = sc_out; - } - - // ---- qk_matmul_output mode 2: post-softcap ---- + // ---- qk_matmul_output mode 2: post-softcap+mask ---- if (has_qk_output && qk_mode == 2) { qk_captured = scores_cur; } diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index ed46060a0b0..3de6ef4108c 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -676,10 +676,8 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Native) { // --------------------------------------------------------------------------- // GQA 3D with softcap — no softcap param in QNN GQA → decomposition. -// DISABLED: accuracy failure on GPU backend -// Max observed delta: ~0.028. -// Tracked as a GPU backend precision issue. -TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Decompose_Softcap) { +// GQA 3D with softcap — no softcap param in QNN GQA → decomposition. +TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { ProviderOptions opts; opts["backend_type"] = "gpu"; opts["offload_graph_io_quantization"] = "0"; From fc22e1b54b2d2505ff92dc444625c58b41d7d082 Mon Sep 17 00:00:00 2001 From: qti-mbadnara Date: Wed, 5 Aug 2026 14:39:27 -0700 Subject: [PATCH 18/20] [QNN EP] Nit Co-authored-by: qti-mattsinc Signed-off-by: qti-mbadnara --- .../providers/qnn/builder/opbuilder/attention_op_builder.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index bb3af4bf330..31c4ea1f2da 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -1294,7 +1294,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn std::string scores_cur = qk_out; // ---- qk_matmul_output mode 0: raw post-QK scores (pre-mask, pre-softcap) ---- - // Per ONNX spec the pipeline is: QK → masks → softcap → softmax + // Per ONNX spec the pipeline is: QK → softcap → masks → softmax // mode 0 = raw QK, mode 1 = post-mask pre-softcap, // mode 2 = post-softcap, mode 3 = post-softmax. std::string qk_captured; // The intermediate captured for qk_matmul_output. From 6e6631f37a0ea6a04d2b52fcddb1269136499d35 Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Thu, 6 Aug 2026 22:33:10 -0700 Subject: [PATCH 19/20] [QNN EP] Switch back to coarse ops --- .../builder/opbuilder/attention_op_builder.cc | 50 +++++++------------ .../test/providers/qnn/attention_test.cc | 4 +- 2 files changed, 18 insertions(+), 36 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 31c4ea1f2da..73da6792108 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -592,11 +592,13 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper&, const OrtNodeUnit&, #endif // SDK version guard for QNN_OP_GROUP_QUERY_ATTENTION // --------------------------------------------------------------------------- -// Helper: emit an ElementWiseBinary (MUL or ADD or DIV) node. +// Helper: emit a dedicated element-wise binary node (ADD, DIVIDE, or MULTIPLY). +// op_type must be one of QNN_OP_ELEMENT_WISE_ADD, QNN_OP_ELEMENT_WISE_DIVIDE, +// or QNN_OP_ELEMENT_WISE_MULTIPLY — no extra params are required for these ops. // --------------------------------------------------------------------------- static Ort::Status AddBinaryOpNode(QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, - uint32_t operation, + const char* op_type, const std::string& lhs_name, const std::string& rhs_name, const std::string& out_name, @@ -614,22 +616,14 @@ static Ort::Status AddBinaryOpNode(QnnModelWrapper& qnn_model_wrapper, const std::string node_name = utils::UniqueNameGenerator().New(node_unit, "_ewb"); - std::vector param_names; - RETURN_IF_ERROR(AddQnnScalar(qnn_model_wrapper, - node_unit.Index(), - node_name, - operation, - QNN_OP_ELEMENT_WISE_BINARY_PARAM_OPERATION, - param_names)); - RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(node_name, QNN_OP_PACKAGE_NAME_QTI_AISW, - QNN_OP_ELEMENT_WISE_BINARY, + op_type, {lhs_name, rhs_name}, {out_name}, - std::move(param_names), + {}, do_op_validation), - "Failed to create ElementWiseBinary node."); + (std::string("Failed to create element-wise binary node: ") + op_type).c_str()); return Ort::Status(); } @@ -757,12 +751,12 @@ static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, // Div(scores, softcap) -> x const std::string div_out = utils::UniqueNameGenerator().New(node_unit, "_softcap_div"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_DIVIDE, + QNN_OP_ELEMENT_WISE_DIVIDE, in_name, sc_name, div_out, shape, dtype, quant_param, /*is_graph_output=*/false, do_op_validation)); - // Tanh(x) -> t [QNN_OP_ELEMENT_WISE_NEURON with TANH param] + // Tanh(x) -> t const std::string tanh_out = utils::UniqueNameGenerator().New(node_unit, "_softcap_tanh"); { QnnTensorWrapper tanh_tensor(tanh_out, QNN_TENSOR_TYPE_NATIVE, dtype, quant_param.Copy(), @@ -770,29 +764,19 @@ static Ort::Status AddSoftcapNode(QnnModelWrapper& qnn_model_wrapper, RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(tanh_tensor)), ("Failed to add softcap Tanh output: " + tanh_out).c_str()); const std::string tanh_node = utils::UniqueNameGenerator().New(node_unit, "_tanh"); - - Qnn_Scalar_t tanh_op_scalar = QNN_SCALAR_INIT; - tanh_op_scalar.dataType = QNN_DATATYPE_UINT_32; - tanh_op_scalar.uint32Value = QNN_OP_ELEMENT_WISE_NEURON_OPERATION_TANH; - QnnParamWrapper tanh_op_param(node_unit.Index(), tanh_node, - QNN_OP_ELEMENT_WISE_NEURON_PARAM_OPERATION, - tanh_op_scalar); - std::vector tanh_param_names = {tanh_op_param.GetParamTensorName()}; - RETURN_IF_NOT(qnn_model_wrapper.AddParamWrapper(std::move(tanh_op_param)), - "Failed to add softcap Tanh operation param."); RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(tanh_node, QNN_OP_PACKAGE_NAME_QTI_AISW, - QNN_OP_ELEMENT_WISE_NEURON, + QNN_OP_TANH, {div_out}, {tanh_out}, - std::move(tanh_param_names), + {}, do_op_validation), "Failed to create softcap Tanh node."); } // Mul(t, softcap) -> out RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + QNN_OP_ELEMENT_WISE_MULTIPLY, tanh_out, sc_name, out_name, shape, dtype, quant_param, /*is_graph_output=*/false, do_op_validation)); @@ -1106,7 +1090,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // current keys are treated uniformly. Without cache, scale K immediately. const std::string q_scaled = utils::UniqueNameGenerator().New(node_unit, "_q_scaled"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + QNN_OP_ELEMENT_WISE_MULTIPLY, q_cur, sqrt_scale_name, q_scaled, q_info.shape, dtype, q_quant, /*is_graph_output=*/false, do_op_validation)); @@ -1116,7 +1100,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // No KV cache: scale K now (standard path). const std::string k_scaled = utils::UniqueNameGenerator().New(node_unit, "_k_scaled"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + QNN_OP_ELEMENT_WISE_MULTIPLY, k_cur, sqrt_scale_name, k_scaled, k_info.shape, dtype, k_quant, /*is_graph_output=*/false, do_op_validation)); @@ -1212,7 +1196,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // This ensures past and current keys are treated uniformly. const std::string k_present_scaled = utils::UniqueNameGenerator().New(node_unit, "_k_present_scaled"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_MULTIPLY, + QNN_OP_ELEMENT_WISE_MULTIPLY, k_cur, sqrt_scale_name, k_present_scaled, k_present_shape, dtype, k_quant, /*is_graph_output=*/false, do_op_validation)); @@ -1375,7 +1359,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const std::string masked_out = utils::UniqueNameGenerator().New(node_unit, "_causal_masked"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_ADD, + QNN_OP_ELEMENT_WISE_ADD, scores_cur, causal_mask_name, masked_out, qk_shape, dtype, q_quant, /*is_graph_output=*/false, do_op_validation)); @@ -1387,7 +1371,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const std::string attn_masked_out = utils::UniqueNameGenerator().New(node_unit, "_attn_masked"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, - QNN_OP_ELEMENT_WISE_BINARY_OPERATION_ADD, + QNN_OP_ELEMENT_WISE_ADD, scores_cur, attn_mask_in, attn_masked_out, qk_shape, dtype, q_quant, /*is_graph_output=*/false, do_op_validation)); diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 3de6ef4108c..4da0433b04e 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -440,9 +440,7 @@ TEST_F(QnnHTPBackendTests, Attention_MQA_3D) { // =========================================================================== // Softcap: scores = softcap * tanh(scores / softcap) -// Gated to real ARM64 hardware: the softcap chain (Div + ElementWiseNeuron -// TANH + Mul) triggers QNN_COMMON_ERROR_MEM_ALLOC during HTP graph -// finalization on the x86_64 HTP simulator. +// Gated to real ARM64 hardware // =========================================================================== #if defined(__aarch64__) || defined(_M_ARM64) From d1d87d16100eb9fb9e63efca7368194667ada30a Mon Sep 17 00:00:00 2001 From: Badri Narayanan Date: Mon, 10 Aug 2026 18:19:54 -0700 Subject: [PATCH 20/20] [QNN EP] Act on review comments --- .../builder/opbuilder/attention_op_builder.cc | 266 ++++++++++-------- .../test/providers/qnn/attention_test.cc | 13 +- 2 files changed, 156 insertions(+), 123 deletions(-) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc index 73da6792108..b3e24afb4b1 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/attention_op_builder.cc @@ -4,7 +4,6 @@ #include #include #include -#include #include #include "QnnOpDef.h" @@ -88,16 +87,25 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper "Attention: past_key present but present_key output not provided"); // ---- nonpad_kv_seqlen (input[6], opset 24) ---- + // The decomposition does not implement the padding-mask computation that + // nonpad_kv_seqlen drives (see ONNX spec defs.cc). Reject it entirely rather + // than silently ignoring it. const bool has_nonpad_kv_seqlen = (num_inputs > 6 && inputs[6].Exists()); - RETURN_IF(has_nonpad_kv_seqlen && has_past_key, - "Attention: nonpad_kv_seqlen and past_key are mutually exclusive per ONNX spec"); + RETURN_IF(has_nonpad_kv_seqlen, + "Attention: nonpad_kv_seqlen is not supported by QNN EP"); + + // ---- softmax_precision ---- + // Cross-dtype cast for softmax accumulation is not implemented. + RETURN_IF(node_helper.HasAttr("softmax_precision"), + "Attention: softmax_precision is not supported by QNN EP"); // ---- qk_matmul_output_mode ---- const int64_t qk_mode = node_helper.Get("qk_matmul_output_mode", static_cast(0)); RETURN_IF(qk_mode < 0 || qk_mode > 3, "Attention: qk_matmul_output_mode must be in [0,3]"); - // If mode != 0 but output[3] is not provided, we just ignore (not an error). - // If output[3] is provided but mode == 0, it's mode-0 (post QK matmul). + // output[3] carries the intermediate tensor selected by qk_matmul_output_mode. + // If output[3] is absent the capture is silently skipped; mode values outside + // [0,3] are rejected above. // ---- scale must be positive if provided ---- if (node_helper.HasAttr("scale")) { @@ -112,6 +120,14 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper RETURN_IF(q_rank != 3 && q_rank != 4, "Attention: Q input must be rank 3 ([B,S_q,n_q*hs]) or rank 4 ([B,n_q,S_q,hs])"); + // ---- Dtype gate: only float16 and float32 are supported ---- + // The scalar buffer helpers (causal mask, sqrt_scale, softcap) write sizeof(float) or + // sizeof(uint16_t) per element. Allowing double (FLOAT_64, 8 bytes) or bfloat16 + // (BFLOAT_16, 2 bytes with a different bit pattern) would silently corrupt those buffers. + RETURN_IF(q_info.qnn_data_type != QNN_DATATYPE_FLOAT_32 && + q_info.qnn_data_type != QNN_DATATYPE_FLOAT_16, + "Attention: only float32 and float16 dtypes are supported by QNN EP"); + TensorInfo k_info{}; RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[1], k_info)); @@ -151,6 +167,18 @@ Ort::Status AttentionOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper "Attention: past_key must be rank 4 ([B,n,S_past,hs])"); } + // ---- attn_mask dtype check ---- + // The ONNX spec converts bool masks via Where(mask, 0, -inf) before adding. + // This builder emits a raw ADD, so a bool tensor would shift logits by 0/1 + // instead of 0/-inf, silently computing wrong attention weights. + const bool has_attn_mask = (num_inputs > 3 && inputs[3].Exists()); + if (has_attn_mask) { + TensorInfo attn_mask_info{}; + RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(inputs[3], attn_mask_info)); + RETURN_IF(attn_mask_info.qnn_data_type == QNN_DATATYPE_BOOL_8, + "Attention: boolean attn_mask is not supported; pre-convert to a float additive bias"); + } + // ---- Full validation: build decomposed nodes with do_op_validation=true ---- std::vector input_names; RETURN_IF_ERROR(ProcessInputs(qnn_model_wrapper, node_unit, logger, input_names, true)); @@ -167,8 +195,7 @@ static bool ShouldUseNativeGQA(QnnBackendType backend, bool has_present_key); // --------------------------------------------------------------------------- -// ProcessInputs — register Q, K, V, attn_mask, past_key, past_value, -// nonpad_kv_seqlen +// ProcessInputs — register Q, K, V, attn_mask, past_key, past_value // --------------------------------------------------------------------------- Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper, const OrtNodeUnit& node_unit, @@ -220,13 +247,11 @@ Ort::Status AttentionOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[4], logger, input_names)); } // input[5] = past_value (optional, KV cache) + // Gated on has_past_key: IsOpSupported enforces both-or-neither for past_key + // and past_value, so past_value is always present whenever past_key is. if (onnx_inputs.size() > 5 && onnx_inputs[5].Exists()) { RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[5], logger, input_names)); } - // input[6] = nonpad_kv_seqlen (optional, opset 24) - if (onnx_inputs.size() > 6 && onnx_inputs[6].Exists()) { - RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, onnx_inputs[6], logger, input_names)); - } return Ort::Status(); } @@ -391,14 +416,9 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, const Qnn_DataType_t dtype = q_info.qnn_data_type; // Derive tensor dimensions needed for 4D output transforms. - uint32_t B = 0, S_q = 0; - if (is_4d) { - B = q_info.shape[0]; - S_q = q_info.shape[2]; - } else { - B = q_info.shape[0]; - S_q = q_info.shape[1]; - } + // B is always shape[0]; S_q depends on layout. + const uint32_t B = q_info.shape[0]; + const uint32_t S_q = is_4d ? q_info.shape[2] : q_info.shape[1]; // ---- Params ---- // For 4D inputs [B, n_q, S_q, hs] the head counts are implicit in the shape; @@ -503,12 +523,10 @@ static Ort::Status EmitNativeGQANode(QnnModelWrapper& qnn_model_wrapper, output_names.push_back(gqa_y_name); // present_key and present_value (output slots 1-2). - // QNN GPU GQA requires real (non-null) tensors for all 3 output slots. - // When the ONNX model does not declare present_key/present_value, allocate - // private NATIVE scratch tensors so QNN can write to them — the caller - // never reads these, but the QNN op def for GQA requires them. - // present_key and present_value (output slots 1-2). - // ShouldUseNativeGQA requires has_present_key=true so both are always declared. + // ShouldUseNativeGQA requires has_present_key=true, so slot 1 is always + // a declared ONNX output. Slot 2 (present_value) follows the same gate + // (IsOpSupported enforces both-or-neither), but the loop below handles + // the absent case with a null tensor for robustness. for (size_t i = 1; i <= 2; ++i) { const bool declared = (onnx_outputs.size() > i && onnx_outputs[i].Exists()); if (declared) { @@ -581,12 +599,17 @@ Ort::Status AttentionOpBuilder::ProcessInputsNativeGQA(const AttentionOpBuilder& const OrtNodeUnit&, const Ort::Logger&, std::vector&) { - return Ort::Status(); + // ShouldUseNativeGQA always returns false on SDK < 2.12, so this is unreachable. + // Return an error rather than silent success to catch any future regression. + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "ProcessInputsNativeGQA: QNN SDK < 2.12, native GQA unavailable"); } static Ort::Status EmitNativeGQANode(QnnModelWrapper&, const OrtNodeUnit&, std::vector&&, bool) { - return Ort::Status(); + // ShouldUseNativeGQA always returns false on SDK < 2.12, so this is unreachable. + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "EmitNativeGQANode: QNN SDK < 2.12, native GQA unavailable"); } #endif // SDK version guard for QNN_OP_GROUP_QUERY_ATTENTION @@ -803,7 +826,7 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, Qnn_DataType_t dtype, const QnnQuantParamsWrapper& quant_param, uint32_t head_ratio, - bool /*do_op_validation*/) { + bool do_op_validation) { // 4D-only GQA expansion that avoids 5D tensors (HTP finalization fails with 5D). // // Goal: produce K_expanded[b, kv*head_ratio+r, s, h] = K[b, kv, s, h] @@ -829,7 +852,7 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(in_name, r1_name, in_shape, r1_shape, dtype, quant_param, - /*do_op_validation=*/false, + do_op_validation, /*is_for_input=*/false)); // Step 2: Tile [1, head_ratio, 1, 1] → [B, head_ratio, n_kv, S*hs] @@ -857,7 +880,7 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, {r1_name}, {tiled_name}, std::move(tile_params), - /*do_op_validation=*/false), + do_op_validation), "Failed to create GQA Tile node."); } @@ -870,14 +893,14 @@ static Ort::Status AddGQAExpandNode(QnnModelWrapper& qnn_model_wrapper, {0u, 2u, 1u, 3u}, tr_shape, dtype, quant_param, - /*do_op_validation=*/false, + do_op_validation, /*is_for_input=*/false)); // Step 4: Reshape [B, n_kv, head_ratio, S*hs] → [B, n_q, S, hs] RETURN_IF_ERROR(qnn_model_wrapper.AddReshapeNode(tr_name, out_name, tr_shape, out_shape, dtype, quant_param, - /*do_op_validation=*/false, + do_op_validation, /*is_for_input=*/false)); return Ort::Status(); } @@ -1040,13 +1063,16 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn size_t opt_idx = 3; const std::string attn_mask_in = has_attn_mask ? input_names[opt_idx++] : std::string{}; const std::string past_key_in = has_past_key ? input_names[opt_idx++] : std::string{}; + // past_value_in gated on has_past_key: IsOpSupported rejects models where + // past_key and past_value are not both-or-neither, so past_value is always + // present whenever past_key is. Relaxing that constraint would corrupt opt_idx. const std::string past_value_in = has_past_key ? input_names[opt_idx++] : std::string{}; std::string q_cur = q_in; std::string k_cur = k_in; std::string v_cur = v_in; - // ---- Step 1: Create a scalar initializer for sqrt(scale) ---- + // ---- Compute sqrt(scale) for Q/K scaling ---- const float scale_default = 1.0f / std::sqrt(static_cast(hs)); const float scale_attr = node_helper.Get("scale", scale_default); const float sqrt_scale = std::sqrt(scale_attr); @@ -1084,10 +1110,9 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn "Failed to add sqrt_scale tensor."); } - // ---- Steps 2-3: Scale Q (always) and K (deferred when KV cache is active) ---- - // For KV cache: present_key = Concat(past_key_raw, K_raw) — store UNSCALED K. - // Scaling of the full K_present happens AFTER the concat so that past and - // current keys are treated uniformly. Without cache, scale K immediately. + // ---- Scale Q (always) and K (deferred to after KV concat when cache is active) ---- + // For KV cache: store unscaled K for concat; scale after concat so past and + // current keys are treated uniformly. Without cache, scale K immediately. const std::string q_scaled = utils::UniqueNameGenerator().New(node_unit, "_q_scaled"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, QNN_OP_ELEMENT_WISE_MULTIPLY, @@ -1109,7 +1134,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // If has_past_key: k_cur is still the raw (unscaled) K here. // Scaling of k_present happens below, after the Concat. - // ---- Steps 4-7 (3D only): Reshape + Transpose Q and K into [B, n, S, hs] ---- + // ---- 3D only: Reshape + Transpose Q and K from BSH into BNSH ---- if (!is_4d) { // Reshape Q: [B, S_q, n_q*hs] -> [B, S_q, n_q, hs] const std::string q_reshaped = utils::UniqueNameGenerator().New(node_unit, "_q_reshaped"); @@ -1157,25 +1182,14 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn } // k_cur is now [B, n_kv, S_k, hs] in BNSH layout. - // ---- GQA expansion of K (if GQA) ---- - if (is_gqa) { - const std::vector k_in_shape = {B, n_kv, S_k, hs}; - const std::vector k_out_shape = {B, n_q, S_k, hs}; - const std::string k_expanded = utils::UniqueNameGenerator().New(node_unit, "_k_gqa"); - RETURN_IF_ERROR(AddGQAExpandNode(qnn_model_wrapper, node_unit, - k_cur, k_expanded, - k_in_shape, k_out_shape, - dtype, k_quant, - head_ratio, - do_op_validation)); - k_cur = k_expanded; - } - // ---- KV cache concat for K ---- + // Concat happens on n_kv-headed tensors (per spec: past_key has kv_num_heads). + // GQA expansion to n_q heads happens AFTER the concat so that the present_key + // output retains kv_num_heads as the spec requires. std::string k_present_name; if (has_past_key) { const uint32_t S_k_total = S_past + S_k; - const std::vector k_present_shape = {B, n_q, S_k_total, hs}; + const std::vector k_present_shape = {B, n_kv, S_k_total, hs}; // If present_key is a graph output, emit it as APP_READ directly from concat. const bool pk_is_graph_out = (onnx_outputs.size() > 1 && onnx_outputs[1].Exists() && @@ -1191,21 +1205,40 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn do_op_validation)); k_cur = k_present_name; S_k = S_k_total; // Update S_k to reflect the full sequence after concat. + } + + // ---- GQA expansion of K (if GQA) ---- + // Runs after KV concat so the expansion covers past+current keys and the + // present_key output retains kv_num_heads as required by the ONNX spec. + if (is_gqa) { + const uint32_t S_k_cur = S_k; // S_k already updated if has_past_key. + const std::vector k_in_shape = {B, n_kv, S_k_cur, hs}; + const std::vector k_out_shape = {B, n_q, S_k_cur, hs}; + const std::string k_expanded = utils::UniqueNameGenerator().New(node_unit, "_k_gqa"); + RETURN_IF_ERROR(AddGQAExpandNode(qnn_model_wrapper, node_unit, + k_cur, k_expanded, + k_in_shape, k_out_shape, + dtype, k_quant, + head_ratio, + do_op_validation)); + k_cur = k_expanded; + } - // Scale the full K_present (past + current) by sqrt(scale) for attention. - // This ensures past and current keys are treated uniformly. + // ---- Scale K (after GQA expansion covers the full sequence) ---- + if (has_past_key) { + // Scale the full K (past + current, already GQA-expanded) by sqrt(scale). + // Deferring until here ensures past and current keys are treated uniformly. + const std::vector k_scaled_shape = {B, n_q, S_k, hs}; const std::string k_present_scaled = utils::UniqueNameGenerator().New(node_unit, "_k_present_scaled"); RETURN_IF_ERROR(AddBinaryOpNode(qnn_model_wrapper, node_unit, QNN_OP_ELEMENT_WISE_MULTIPLY, k_cur, sqrt_scale_name, k_present_scaled, - k_present_shape, dtype, k_quant, + k_scaled_shape, dtype, k_quant, /*is_graph_output=*/false, do_op_validation)); k_cur = k_present_scaled; } - // ---- Steps 12-13 (3D only): Reshape + Transpose V into [B, n_kv, S_k_cur, v_hs] ---- - // NOTE: S_k used here is the original S_k of V (before KV concat). - // We do V transform BEFORE KV concat so V uses original S_k dimensions. + // ---- 3D only: Reshape + Transpose V from BSH into BNSH ---- const uint32_t S_k_orig = has_past_key ? (S_k - S_past) : S_k; std::string v_cur_4d = v_cur; @@ -1234,25 +1267,12 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn } // v_cur_4d is now [B, n_kv, S_k_orig, v_hs] in BNSH layout. - // ---- GQA expansion of V (if GQA) ---- - if (is_gqa) { - const std::vector v_in_shape = {B, n_kv, S_k_orig, v_hs}; - const std::vector v_out_shape = {B, n_q, S_k_orig, v_hs}; - const std::string v_expanded = utils::UniqueNameGenerator().New(node_unit, "_v_gqa"); - RETURN_IF_ERROR(AddGQAExpandNode(qnn_model_wrapper, node_unit, - v_cur_4d, v_expanded, - v_in_shape, v_out_shape, - dtype, v_quant, - head_ratio, - do_op_validation)); - v_cur_4d = v_expanded; - } - // ---- KV cache concat for V ---- + // Same ordering as K: concat on n_kv-headed tensors first, then GQA-expand. std::string v_present_name; if (has_past_key) { const uint32_t S_k_total = S_k; // already updated above. - const std::vector v_present_shape = {B, n_q, S_k_total, v_hs}; + const std::vector v_present_shape = {B, n_kv, S_k_total, v_hs}; const bool pv_is_graph_out = (onnx_outputs.size() > 2 && onnx_outputs[2].Exists() && qnn_model_wrapper.IsGraphOutput(onnx_outputs[2].name)); @@ -1268,7 +1288,23 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn v_cur_4d = v_present_name; } - // ---- Step 8: MatMul Q * K^T (transpose_in1) -> [B, n_q, S_q, S_k] ---- + // ---- GQA expansion of V (if GQA) ---- + // Runs after KV concat so present_value retains kv_num_heads per the ONNX spec. + if (is_gqa) { + const uint32_t S_v_cur = S_k; // S_k updated to S_total when has_past_key. + const std::vector v_in_shape = {B, n_kv, S_v_cur, v_hs}; + const std::vector v_out_shape = {B, n_q, S_v_cur, v_hs}; + const std::string v_expanded = utils::UniqueNameGenerator().New(node_unit, "_v_gqa"); + RETURN_IF_ERROR(AddGQAExpandNode(qnn_model_wrapper, node_unit, + v_cur_4d, v_expanded, + v_in_shape, v_out_shape, + dtype, v_quant, + head_ratio, + do_op_validation)); + v_cur_4d = v_expanded; + } + + // ---- MatMul Q * K^T -> [B, n_q, S_q, S_k] ---- const std::string qk_out = utils::UniqueNameGenerator().New(node_unit, "_qk_out"); const std::vector qk_shape = {B, n_q, S_q, S_k}; RETURN_IF_ERROR(AddMatMulNode(qnn_model_wrapper, node_unit, @@ -1277,19 +1313,28 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn /*transpose_in1=*/true, do_op_validation)); std::string scores_cur = qk_out; - // ---- qk_matmul_output mode 0: raw post-QK scores (pre-mask, pre-softcap) ---- - // Per ONNX spec the pipeline is: QK → softcap → masks → softmax - // mode 0 = raw QK, mode 1 = post-mask pre-softcap, - // mode 2 = post-softcap, mode 3 = post-softmax. + // ---- qk_matmul_output mode 0: raw post-QK scores (pre-softcap, pre-mask) ---- + // Per the published ONNX 1.23 spec (1.23 errata reversed the earlier ordering): + // mode 0 = raw QK (pre-softcap, pre-mask) + // mode 1 = post-softcap, pre-mask + // mode 2 = post-softcap+mask (pre-softmax) + // mode 3 = post-softmax + // NOTE: cmake/external/onnx is pinned at v1.20.1 (pre-errata) where mode 1 + // was post-mask, pre-softcap — the opposite. The code below is correct per the + // published spec; do not "fix" it by grepping the vendored submodule. std::string qk_captured; // The intermediate captured for qk_matmul_output. if (has_qk_output && qk_mode == 0) { qk_captured = scores_cur; } - // ---- Softcap BEFORE masks (per ONNX spec) ---- + // ---- Softcap BEFORE masks (per published ONNX spec) ---- // From the ONNX spec function body: MatMul → softcap → attn_mask+Add → Softmax. - // Softcap must come first: applying the mask (-1e9) before softcap would clamp - // it to -softcap (e.g. -5.0), making masked positions visible to softmax. + // Softcap must come first: applying the causal mask (-1e9) before softcap would + // clamp it to -softcap (e.g. -5.0), making masked positions visible to softmax. + // + // NOTE: cmake/external/onnx is pinned at v1.20.1 (pre-1.23 errata), where the + // submodule shows mask-before-softcap. The published 1.23 errata reverses this. + // The ordering here is correct per the published spec. if (softcap != 0.0f) { const std::string sc_out = utils::UniqueNameGenerator().New(node_unit, "_scores_softcap"); RETURN_IF_ERROR(AddSoftcapNode(qnn_model_wrapper, node_unit, @@ -1307,13 +1352,21 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn // ---- Causal mask (is_causal=1): ADD static lower-triangular mask ---- // With KV cache: offset = S_past so that row i attends to positions <= i+S_past. + // + // Shape is [1, 1, S_q, S_k] — QNN broadcasts this over [B, n_q, S_q, S_k] in ADD. + // All batch and head positions share the same lower-triangular pattern, so there + // is no need to materialize the full [B, n_q, S_q, S_k] tensor. + // + // -1e9f (fp32) / -1e4f (fp16) rather than -inf: HTP V73 HVX does not reliably + // propagate IEEE-754 -inf through ADD. Using a large finite negative ensures + // softmax produces ~0 for masked positions without NaN in the output row. if (is_causal != 0) { const uint32_t offset = S_past; // 0 for no KV cache path. const std::string causal_mask_name = utils::UniqueNameGenerator().New(node_unit, "_causal_mask"); { - std::vector mask_shape = {B, n_q, S_q, S_k}; - const size_t total = static_cast(B) * static_cast(n_q) * - static_cast(S_q) * static_cast(S_k); + // Broadcast shape: [1, 1, S_q, S_k] + const std::vector mask_shape = {1u, 1u, S_q, S_k}; + const size_t total = static_cast(S_q) * static_cast(S_k); std::vector mask_bytes; if (dtype == QNN_DATATYPE_FLOAT_16) { @@ -1321,28 +1374,18 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn const uint16_t neg_raw = fp16_large_neg.val; mask_bytes.resize(total * sizeof(uint16_t)); uint16_t* mask_ptr = reinterpret_cast(mask_bytes.data()); - for (uint32_t b = 0; b < B; ++b) { - for (uint32_t h = 0; h < n_q; ++h) { - for (uint32_t i = 0; i < S_q; ++i) { - for (uint32_t j = 0; j < S_k; ++j) { - const size_t idx = ((static_cast(b) * n_q + h) * S_q + i) * S_k + j; - mask_ptr[idx] = (j <= i + offset) ? static_cast(0u) : neg_raw; - } - } + for (uint32_t i = 0; i < S_q; ++i) { + for (uint32_t j = 0; j < S_k; ++j) { + mask_ptr[i * S_k + j] = (j <= i + offset) ? static_cast(0u) : neg_raw; } } } else { constexpr float large_neg = -1e9f; mask_bytes.resize(total * sizeof(float)); float* mask_ptr = reinterpret_cast(mask_bytes.data()); - for (uint32_t b = 0; b < B; ++b) { - for (uint32_t h = 0; h < n_q; ++h) { - for (uint32_t i = 0; i < S_q; ++i) { - for (uint32_t j = 0; j < S_k; ++j) { - const size_t idx = ((static_cast(b) * n_q + h) * S_q + i) * S_k + j; - mask_ptr[idx] = (j <= i + offset) ? 0.0f : large_neg; - } - } + for (uint32_t i = 0; i < S_q; ++i) { + for (uint32_t j = 0; j < S_k; ++j) { + mask_ptr[i * S_k + j] = (j <= i + offset) ? 0.0f : large_neg; } } } @@ -1351,7 +1394,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn QNN_TENSOR_TYPE_STATIC, dtype, QnnQuantParamsWrapper{}, - std::move(mask_shape), + std::vector(mask_shape), std::move(mask_bytes)); RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(mask_tensor)), "Failed to add causal mask tensor."); @@ -1383,7 +1426,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn qk_captured = scores_cur; } - // ---- Step 11: Softmax (axis=3) -> [B, n_q, S_q, S_k] ---- + // ---- Softmax (axis=3) -> [B, n_q, S_q, S_k] ---- const std::string softmax_out = utils::UniqueNameGenerator().New(node_unit, "_softmax_out"); RETURN_IF_ERROR(AddSoftmaxNode(qnn_model_wrapper, node_unit, scores_cur, softmax_out, @@ -1396,7 +1439,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn qk_captured = attn_weights; } - // ---- Step 14: MatMul attn_weights * V -> [B, n_q, S_q, v_hs] ---- + // ---- MatMul attn_weights * V -> [B, n_q, S_q, v_hs] ---- const std::string y_pre_transpose = utils::UniqueNameGenerator().New(node_unit, "_y_pre_transpose"); const std::vector y_pre_shape = {B, n_q, S_q, v_hs}; @@ -1405,7 +1448,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn y_pre_shape, dtype, q_quant, /*transpose_in1=*/false, do_op_validation)); - // ---- Steps 15-16 (3D outputs): Transpose + Reshape back to [B, S_q, n_q*v_hs] ---- + // ---- 3D outputs: Transpose + Reshape Y back to BSH [B, S_q, n_q*v_hs] ---- const std::string& final_output_name = onnx_outputs[0].name; const bool is_graph_output = qnn_model_wrapper.IsGraphOutput(final_output_name); TensorInfo y_info{}; @@ -1468,20 +1511,19 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn } // ---- Register KV cache outputs (if not already APP_READ) ---- - // When has_past_key is true and present_key/present_value are expected outputs, - // the concat nodes above already produced them. If they are graph outputs they - // were created as APP_READ. If they are not graph outputs but the ONNX node - // declares them as outputs we still need to register them. + // The concat nodes produced present_key [B,n_kv,S_k,hs] and present_value + // [B,n_kv,S_k,v_hs] as NATIVE or APP_READ tensors. If an ONNX output slot + // is declared but the tensor was created as a NATIVE intermediate (i.e. the + // name differs), expose it with an identity Reshape. if (has_past_key) { // present_key = output[1] if (onnx_outputs.size() > 1 && onnx_outputs[1].Exists()) { const bool is_go = qnn_model_wrapper.IsGraphOutput(onnx_outputs[1].name); if (!is_go && onnx_outputs[1].name != k_present_name) { - // Need to expose it. The concat output already has the right shape. RETURN_IF_ERROR(RegisterIntermediateAsOutput(qnn_model_wrapper, node_unit, k_present_name, onnx_outputs[1].name, - {B, n_q, S_k, hs}, + {B, n_kv, S_k, hs}, dtype, k_quant, do_op_validation)); } @@ -1493,7 +1535,7 @@ Ort::Status AttentionOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn RETURN_IF_ERROR(RegisterIntermediateAsOutput(qnn_model_wrapper, node_unit, v_present_name, onnx_outputs[2].name, - {B, n_q, S_k, v_hs}, + {B, n_kv, S_k, v_hs}, dtype, v_quant, do_op_validation)); } diff --git a/onnxruntime/test/providers/qnn/attention_test.cc b/onnxruntime/test/providers/qnn/attention_test.cc index 4da0433b04e..139665985af 100644 --- a/onnxruntime/test/providers/qnn/attention_test.cc +++ b/onnxruntime/test/providers/qnn/attention_test.cc @@ -506,7 +506,6 @@ TEST_F(QnnHTPBackendTests, Attention_KVCache_4D) { TestInputDef({1, 4, 4, 16}, true, -1.0f, 1.0f), // past_value (initializer) {test::MakeAttribute("is_causal", static_cast(0))}), opts, /*opset_version=*/24, - // KV cache (Concat) adds extra operations; fp16 rounding accumulates more error. EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } @@ -566,7 +565,7 @@ TEST_F(QnnHTPBackendTests, Attention_DebugOutput_Mode3) { #if defined(_M_ARM64) // --------------------------------------------------------------------------- -// Native GQA path +// Native GQA path (GPU + causal + no softcap/attn_mask/qk_output) // --------------------------------------------------------------------------- // GQA 3D with KV cache, causal — native QNN_OP_GROUP_QUERY_ATTENTION. @@ -592,10 +591,6 @@ TEST_F(QnnGPUBackendTests, DISABLED_Attention_GPU_GQA_3D_Native_KVCache) { EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(2e-3f)}); } -// --------------------------------------------------------------------------- -// Native GQA path (GPU + causal + no softcap/attn_mask/qk_output) -// --------------------------------------------------------------------------- - // GQA 3D, head_ratio=2, causal, no KV cache → native QNN_OP_GROUP_QUERY_ATTENTION. // Q [1,8,32] (4 heads × 8), K/V [1,8,16] (2 heads × 8). TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Native) { @@ -652,8 +647,6 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_MHA_3D_Native) { } // GQA 4D BNSH, causal, no KV cache → native (EmitNativeGQANode inserts Transpose+Reshape). -// DISABLED: after 4D→3D transforms the QNN node has unpacked K/V inputs, which is not -// supported by the GPU backend in QAIRT 2.48; re-enable in 2.50. TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Native) { ProviderOptions opts; opts["backend_type"] = "gpu"; @@ -673,7 +666,6 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Native) { // Decomposition path on GPU — each test disqualifies one native condition. // --------------------------------------------------------------------------- -// GQA 3D with softcap — no softcap param in QNN GQA → decomposition. // GQA 3D with softcap — no softcap param in QNN GQA → decomposition. TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { ProviderOptions opts; @@ -693,8 +685,7 @@ TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_3D_Decompose_Softcap) { EPVerificationParams{ExpectedEPNodeAssignment::All, ElementwiseAbsoluteVerifier(1e-3f)}); } -// GQA 4D BNSH inputs, no KV cache — no past_key → decomposition. -// GQA 4D BNSH inputs, no KV cache → decomposition (is_causal=0 disqualifies native path). +// GQA 4D BNSH inputs, is_causal=0 → decomposition (causal required for native path). // Q [1,4,8,8] (n_q=4), K/V [1,2,8,8] (n_kv=2). TEST_F(QnnGPUBackendTests, Attention_GPU_GQA_4D_Decompose) { ProviderOptions opts;