[QNN EP] Add support for Attention Op

[qti-mbadnara]

Description

Adds AttentionOpBuilder to the QNN EP, enabling the standard ONNX ai.onnx::Attention operator (opset 23 and 24) to run on the Hexagon HTP and Adreno GPU backends. The builder decomposes the fused Attention node into elementary QNN ops for HTP, and for the GPU backend routes to the native QNN_OP_GROUP_QUERY_ATTENTION kernel when possible.

---

Motivation & Context

The QNN EP had no registration for ai.onnx::Attention. Every Attention node fell back to CPU EP; with session.disable_cpu_ep_fallback=1 session creation failed outright:

FAIL : This session contains graph nodes that are assigned to the default CPU EP,
but fallback to CPU EP has been explicitly disabled by the user.

---

Routing

The builder uses a two-track strategy based on backend and node attributes:

Condition	Path
GPU backend + n_q % n_kv == 0 + is_causal=1 + no softcap/attn_mask/qk_output + present_key declared	Native QNN_OP_GROUP_QUERY_ATTENTION
All other cases (HTP, non-causal, softcap, attn_mask, no present_key)	Decomposition into elementary QNN ops

The native path covers both MHA (n_q == n_kv) and GQA/MQA (n_q > n_kv). Both 3D BSH and 4D BNSH inputs are supported: for 4D, Transpose+Reshape nodes are inserted before Q/K/V and a Reshape+Transpose is inserted after Y to bridge the layout.

Why decomposition for HTP: QNN SDK ≤ 2.49 has no MHA primitive for HTP. QNN_OP_GROUP_QUERY_ATTENTION is GPU-only (Adreno). The ONNX Attention spec defines a reference body of Scale + MatMul + Softmax + MatMul; this PR implements that body as QNN node sequences.

Why decomposition on GPU for some cases: The native QNN_OP_GROUP_QUERY_ATTENTION does not support softcap or an additive attn_mask input, and is always causal. Cases that need these features fall to the decomposition path, which runs on GPU using individual QNN ops.

---

GPU Native Path (QNN_OP_GROUP_QUERY_ATTENTION)

A single node is emitted with:

• Params: NUM_HEADS, KV_NUM_HEADS, SCALE, DO_ROTARY=0
• Inputs: query (BSH), synthetic seqlens_k (INT32 [B] = S_past + S_k − 1), synthetic total_seq_len (INT32 0D scalar = S_past + S_k), key, value, past_key/past_value (4D BNSH, passed directly)
• Outputs: Y (BSH), present_key / present_value (4D BNSH, when KV cache is active)

seqlens_k and total_seq_len are synthesized from static input shapes; they are not present in the ONNX ai.onnx::Attention signature but required by QNN GQA.

For 4D BNSH inputs, the builder transparently inserts Transpose(0,2,1,3)+Reshape before Q/K/V (converting [B,n,S,hs] → [B,S,n*hs]) and Reshape+Transpose(0,2,1,3) after Y (converting [B,S,n*hs] → [B,n,S,hs]). past_key/past_value are already BNSH and are passed straight through.

---

Decomposition

3D inputs [B, S, n·hs] (BSH layout)

Step	QNN Op	Output shape
1	Static sqrt(scale) scalar	—
2	ElementWiseMultiply Q·√s	[B, S_q, n_q·hs]
3	ElementWiseMultiply K·√s (no KV cache)	[B, S_k, n_kv·hs]
4–5	Reshape + Transpose(0,2,1,3) Q → BNSH	[B, n_q, S_q, hs]
6–7	Reshape + Transpose(0,2,1,3) K → BNSH	[B, n_kv, S_k, hs]
8	KV cache: Concat(past_key, K, axis=2) (cache only)	[B, n_kv, S_total, hs]
9	GQA expand K: Reshape→Tile→Transpose→Reshape (GQA only)	[B, n_q, S_k, hs]
10	ElementWiseMultiply K·√s (KV cache path only, after concat+expand)	[B, n_q, S_total, hs]
11	MatMul(Q, K^T, transpose_in1=true)	[B, n_q, S_q, S_k]
12	Softcap: ElementWiseDivide→Tanh→ElementWiseMultiply (if softcap!=0)	[B, n_q, S_q, S_k]
13	Static causal mask ElementWiseAdd (is_causal=1)	[B, n_q, S_q, S_k]
14	User attn_mask ElementWiseAdd (if provided)	[B, n_q, S_q, S_k]
15	Softmax(axis=3)	[B, n_q, S_q, S_k]
16–17	Reshape + Transpose(0,2,1,3) V → BNSH	[B, n_kv, S_k, v_hs]
18	KV cache: Concat(past_value, V, axis=2) (cache only)	[B, n_kv, S_total, v_hs]
19	GQA expand V (GQA only)	[B, n_q, S_k, v_hs]
20	MatMul(attn_weights, V)	[B, n_q, S_q, v_hs]
21–22	Transpose(0,2,1,3) + Reshape Y	[B, S_q, n_q·v_hs]

The qk_matmul_output_mode capture points (per published ONNX 1.23 spec — note that cmake/external/onnx is pinned at v1.20.1 pre-errata which has the opposite ordering):

• mode 0 — raw QK scores (after step 11, before softcap)
• mode 1 — post-softcap, pre-mask (after step 12)
• mode 2 — post-softcap+mask, pre-softmax (after step 14)
• mode 3 — post-softmax / attn_weights (after step 15)

For 4D inputs [B, n, S, hs] (BNSH layout): steps 4–7 and 21–22 are skipped — inputs are already in BNSH format.

Operation order — softcap before masks (per ONNX 1.23 spec): Step 12 applies softcap to the raw QK scores before masks. Applying masks before softcap would clamp -1e9 causal mask values to -softcap, making masked positions visible to softmax.

GQA head expansion (floor-division semantics matching the ONNX function body):

• K[b, kv*ratio+r, s, h] = K_orig[b, kv, s, h] → [K0,K0,...,K1,K1,...]
• 4D approach (avoids 5D HTP limitations): Reshape[B,1,n_kv,S*hs]→Tile[1,ratio,1,1]→Transpose(0,2,1,3)→Reshape[B,n_q,S,hs]
• GQA expansion runs after the KV concat so that present_key/present_value retain kv_num_heads per the ONNX spec.

Causal mask (is_causal=1): Static [1, 1, S_q, S_k] lower-triangular float tensor (broadcast over [B, n_q, S_q, S_k] in ADD), values 0.0 for j ≤ i+offset, -1e9 otherwise. -inf is avoided because HTP V73 HVX does not reliably propagate IEEE-754 -inf through ADD. Offset = S_past when KV cache is active, 0 otherwise.

Softcap (softcap != 0): ElementWiseDivide(scores, sc) → QNN_OP_TANH → ElementWiseMultiply(result, sc).

KV cache ordering: present_key = Concat(past_key, K_raw) at kv_num_heads heads, then GQA-expanded to n_q heads, then scaled by √s. This ensures past and current keys are treated uniformly, and present_key/present_value outputs carry kv_num_heads as the spec requires.

---

Supported Configuration

Feature	HTP	GPU
Input rank	3D BSH or 4D BNSH	3D BSH or 4D BNSH
Opsets	23, 24	23, 24
Dtypes	float32, float16	float32
MHA (n_q == n_kv, is_causal=1)	Decomposition ✓	Native QNN_OP_GROUP_QUERY_ATTENTION ✓
GQA / MQA (n_q > n_kv, is_causal=1)	Decomposition ✓	Native QNN_OP_GROUP_QUERY_ATTENTION ✓
Any attention with is_causal=0	Decomposition ✓	Decomposition ✓
scale	Any positive float ✓	Any positive float ✓
attn_mask (float additive bias)	✓	Decomposition ✓
softcap	via QNN_OP_TANH ✓	via QNN_OP_TANH (decomposition) ✓
KV cache (past_key/present_key)	✓	✓ (native path includes past/present)
qk_matmul_output modes 0–3	✓	Decomposition ✓
Shape	Static only	Static only

---

Unsupported / Rejected

• Dynamic shape dimensions on Q, K, or V
• scale <= 0
• nonpad_kv_seqlen — padding-mask computation not implemented; rejected outright
• softmax_precision — cross-dtype softmax accumulation not implemented; rejected outright
• past_key without present_key (or vice versa)
• Boolean attn_mask — spec requires Where(mask, 0, -inf) conversion; builder emits raw ADD, so bool masks are rejected. Pre-convert to a float additive bias before passing to this EP.
• double and bfloat16 dtypes — only float32 and float16 are supported; other types would silently corrupt scalar buffer helpers (causal mask, sqrt_scale, softcap)

[qti-mattsinc]

@qti-mbadnara, can we map onnx Attention onto QNN_OP_GROUP_QUERY_ATTENTION for the GPU backend?

Why decomposition, not a native QNN op:
QNN SDK ≤ 2.49 has no MHA primitive for HTP. QNN_OP_GROUP_QUERY_ATTENTION is GPU-only (Adreno). The ONNX Attention spec defines a reference body of Scale + MatMul + Softmax + MatMul; this PR implements that body as QNN node sequences.

HTP may not support this, but GPU can. And perf on GPU using the QNN GQA op will be better than the decomposed sequence.

[johnpaultaken]

I agree that the onnx Attention need to map to QNN_OP_GROUP_QUERY_ATTENTION in the normal case. It can resort to this splitting only for the special HTP case temporarily because it would be beneficial for HTP also to implement GQA eventually. I guess the onnx document specifies splitting Attention op to Scale + MatMul + Softmax + MatMul only as a temporary measure until implementations can catch up with the full Attention op kernels. Afaik the GQA is as generic as the new ONNX Attention op. Still we can also add the new ONNX Attention op to QNN if desired. Please let us know your thoughts.

[qti-mattsinc]

Please fix decomposition ordering with softcap

[qti-mattsinc]

LGTM

[qti-yuduo]

Reviewed against the ONNX Attention spec and the QNN SDK headers. Architecture and the two-path split look right, and the SDK 2.12 guard is load-bearing (local QnnOpDef.h is 2.4.0 with no GROUP_QUERY symbols).

The softcap-before-mask order is correct — worth noting the vendored cmake/external/onnx is pinned at v1.20.1, which is pre-1.23-errata and has the opposite order. I initially flagged this as a bug against the submodule before checking the published spec. See the comment on L1282.

Two blocking issues (L1336 dtype buffer, L1370 bool mask). Everything else is minor.

Also: IsOpSupported re-runs the decomposition with do_op_validation=true, so most shape mismatches become a clean CPU fallback rather than a bad compile. The bool-mask issue is the exception — it builds a valid graph that computes the wrong thing.

[qti-mbadnara]

Hi @qti-yuduo thanks for the detailed review! Fixed the blockers.