Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Attention

Kernel card Files
xet
 Community

ai.onnx.Attention

ai.onnx · standard ONNX operator · ONNX opset ≥ 24

Description

Computes scaled dot-product attention as softmax(Q * K^T * scale) * V, with optional masks and softcap. Supports MHA, GQA, and MQA via q_num_heads and kv_num_heads, with 4D (batch, heads, sequence, head_size) or 3D (batch, sequence, hidden) inputs. Rank-4 routes support past/present KV caches and masks over the joined key sequence. This package implements Attention-24, including causal alignment for its internal past_key cache; it does not expose nonpad_kv_seqlen or Attention-25's left_window_size and right_window_size attributes.

See the ONNX Attention spec for the reference semantics.

Inputs

Name Bind key Logical dtype Rank Shape Description Presence
Q q T — — Query tensor; 4D (batch, q_num_heads, q_seq, head_size) or 3D (batch, q_seq, q_hidden_size). required
K k T — — Key tensor; 4D (batch, kv_num_heads, kv_seq, head_size) or 3D (batch, kv_seq, k_hidden_size). required
V v T — — Value tensor; 4D (batch, kv_num_heads, kv_seq, v_head_size) or 3D (batch, kv_seq, v_hidden_size). required
attn_mask attn_mask M — — Optional attention mask broadcastable to (batch, q_num_heads, q_seq, kv_seq); a true boolean permits attention, while a mask of the same floating-point type as Q, K, and V is added to the scores. optional
past_key past_key T 4 — Optional cached keys of shape (batch, kv_num_heads, past_sequence_length, head_size); the cached tokens precede K on the sequence axis. optional
past_value past_value T 4 — Optional cached values of shape (batch, kv_num_heads, past_sequence_length, v_head_size); supplied together with past_key. optional

Outputs

Name Bind key Logical dtype Rank Shape Description Presence
Y y T same as Q derived; see description Attention output; 4D (batch, q_num_heads, q_seq, v_head_size) or 3D (batch, q_seq, hidden_size). required
present_key present_key T 4 derived; see description Optional joined keys, past_key followed by K, of shape (batch, kv_num_heads, past_sequence_length + kv_sequence_length, head_size). optional
present_value present_value T 4 derived; see description Optional joined values, past_value followed by V, of shape (batch, kv_num_heads, past_sequence_length + kv_sequence_length, v_head_size). optional

Attributes

Attributes and default values (overridable per request):

Attribute Default Description
is_causal 0 When set to 1, applies causal masking. Without a cache, query index i attends keys through i; with past_key, it attends keys through past_sequence_length + i, aligning the mask to the query's absolute position in the joined cache.
qk_matmul_output_mode 0 Selects the stage exposed by the optional QK-matmul output. This package does not expose that output yet, so only the default mode 0 is accepted.
softcap 0 If non-zero, attention logits are capped via tanh(logits / softcap) * softcap before the softmax.
scale — Scalar multiplier applied to Q * K^T; defaults to 1 / sqrt(head_size).
q_num_heads — Number of query attention heads; required when Q, K, V are 3D tensors.
kv_num_heads — Number of key/value attention heads; required when Q, K, V are 3D tensors (set equal to q_num_heads for MHA, or 1 for MQA).
softmax_precision — Optional TensorProto element-type code for the softmax computation. Only FLOAT (1) is implemented; FLOAT16 (10) is not yet supported.

Type constraints

Variable Allowed dtypes
T float32, float16
M float32, float16, bool

Device requirements

Some implementation variants require subgroup-matrix, shader-f16, and subgroups. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.

Files

Use with @huggingface/kernels

The loader derives every required output's shape and logical dtype from the manifest contract and this call. It then allocates the result tensors automatically.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model version.

Replace each *Data placeholder with a typed array containing the corresponding input data.

import { getKernel } from "@huggingface/kernels";

const kernel = await getKernel("webgpu-kernels/ai.onnx.Attention", { version: 1 });
const { y } = await kernel({
  q: { data: qData, shape: [1, 1, 1, 2] },
  k: { data: kData, shape: [1, 1, 2, 2] },
  v: { data: vData, shape: [1, 1, 2, 1] },
});
Downloads last month
-
kernel
webgpu
wgsl
apache-2.0
WebGPU

Requires WebGPU support. See the compatibility table.