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

Outputs

Attributes

Attributes and default values (overridable per request):

Type constraints

Implementation variants

One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.

• materialized_sgmat_f32 — Materializes scores and row statistics so score generation and value application can use subgroup-matrix tiles. It applies to unmasked, non-causal float32 attention without softcap when subgroup-matrix support and scratch capacity are available.
• materialized_gemm_f32 — Materializes scores and row statistics for tiled score and value-application passes. It is the portable materialized route for unmasked, non-causal float32 attention without softcap when the intermediates fit.
• rank4_kv_cache_online — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_online_key_only — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_online_value_only — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_online_y_only — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_mask_online — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_mask_online_key_only — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_mask_online_value_only — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_kv_cache_mask_online_y_only — Joins paired past_key/past_value inputs with K/V, retaining an omitted present cache in private intermediate storage when necessary, then runs the online rank-4 attention kernel over the joined tensors. Requested present_key and present_value outputs remain independently optional.
• rank4_present_outputs_online — Copies current K and V to the independently requested present cache outputs when no past cache is supplied, while the online rank-4 attention pass reads K and V directly.
• rank4_present_outputs_mask_online — Copies current K and V to the independently requested present cache outputs when no past cache is supplied, while the online rank-4 attention pass reads K and V directly.
• rank4_present_key_online — Copies the independently requested current key or value tensor to its present cache output when no past cache is supplied, while the online rank-4 attention pass reads K and V directly.
• rank4_present_value_online — Copies the independently requested current key or value tensor to its present cache output when no past cache is supplied, while the online rank-4 attention pass reads K and V directly.
• rank4_present_key_mask_online — Copies the independently requested current key or value tensor to its present cache output when no past cache is supplied, while the online rank-4 attention pass reads K and V directly.
• rank4_present_value_mask_online — Copies the independently requested current key or value tensor to its present cache output when no past cache is supplied, while the online rank-4 attention pass reads K and V directly.
• flash_online_tiled_cluster_nosg — Processes a tile of query rows against successive key tiles with online softmax, sharing K/V staging across the query tile without materializing the score matrix.
• flash_online_tiled_mask_cluster_nosg — Processes a tile of query rows against successive key tiles with online softmax, sharing K/V staging across the query tile without materializing the score matrix.
• flash_online_tiled_cluster — Processes a tile of query rows against successive key tiles with online softmax, sharing K/V staging across the query tile without materializing the score matrix.
• flash_online_tiled_mask_cluster — Processes a tile of query rows against successive key tiles with online softmax, sharing K/V staging across the query tile without materializing the score matrix.
• rank4_no_mask_2pass — Stores unnormalized softmax weights and row denominators, then applies them to V in a tiled second pass. This admits value-head widths beyond the generic online route's workgroup-storage bound.
• rank4_mask_2pass — Stores unnormalized softmax weights and row denominators, then applies them to V in a tiled second pass. This admits value-head widths beyond the generic online route's workgroup-storage bound.
• flash_decode_splitk_nosg — Partitions the key sequence across workgroups to create parallelism when a short-query grid is too small, then merges online-softmax state from scratch. The merge reassociates floating-point reductions relative to a single key sweep.
• flash_online_tiled_decode_splitk — Partitions the key sequence across workgroups to create parallelism when a short-query grid is too small, then merges online-softmax state from scratch. The merge reassociates floating-point reductions relative to a single key sweep.
• flash_short_query_splitk_nosg — Partitions the key sequence across workgroups to create parallelism when a short-query grid is too small, then merges online-softmax state from scratch. The merge reassociates floating-point reductions relative to a single key sweep.
• flash_short_query_splitk — Partitions the key sequence across workgroups to create parallelism when a short-query grid is too small, then merges online-softmax state from scratch. The merge reassociates floating-point reductions relative to a single key sweep.
• flash_masked_decode_splitk — Partitions the key sequence across workgroups to create parallelism when a short-query grid is too small, then merges online-softmax state from scratch. The merge reassociates floating-point reductions relative to a single key sweep.
• flash_masked_short_query_splitk — Partitions the key sequence across workgroups to create parallelism when a short-query grid is too small, then merges online-softmax state from scratch. The merge reassociates floating-point reductions relative to a single key sweep.
• flash_cluster_splitkv — Tiles query rows while partitioning a long key/value sequence across workgroups, increasing parallelism without materializing scores. It merges per-partition online-softmax state from scratch, which reassociates floating-point reductions relative to a single key sweep.
• flash_cluster_splitkv_nosg — Tiles query rows while partitioning a long key/value sequence across workgroups, increasing parallelism without materializing scores. It merges per-partition online-softmax state from scratch, which reassociates floating-point reductions relative to a single key sweep.
• flash_cluster_splitkv_mask — Tiles query rows while partitioning a long key/value sequence across workgroups, increasing parallelism without materializing scores. It merges per-partition online-softmax state from scratch, which reassociates floating-point reductions relative to a single key sweep.
• flash_cluster_splitkv_mask_nosg — Tiles query rows while partitioning a long key/value sequence across workgroups, increasing parallelism without materializing scores. It merges per-partition online-softmax state from scratch, which reassociates floating-point reductions relative to a single key sweep.

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

• metadata.json — kernel metadata (id, digests, per-variant templates, provenance)
• manifest.json — the op contract (source of truth)
• test.json — correctness cases
• bench.json — benchmark + tuning cases
• attention-rank4-apply-tiled.wgsl.jinja
• attention-rank4-online.wgsl.jinja
• attention-rank4-softmax.wgsl.jinja
• attention-rank4-tiled.wgsl.jinja
• attention-rank4.wgsl.jinja
• attn-flash-decode-splitk-merge.wgsl.jinja
• attn-flash-decode-splitk.wgsl.jinja
• attn-flash-online.wgsl.jinja
• attn-flash-prefill-cluster.wgsl.jinja
• attn-flash-q32-broadcast.wgsl.jinja
• attn-kv-cache-concat.wgsl.jinja
• attn-materialized-apply-f32.wgsl.jinja
• attn-materialized-rowstats-combine-f32.wgsl.jinja
• attn-materialized-score-f32.wgsl.jinja
• attn-materialized-sgmat-f32.wgsl.jinja
• attn-online-scalar.wgsl.jinja
• datamove-flat-copy.wgsl.jinja

Use with @huggingface/kernels

npm install --save-exact @huggingface/kernels@0.0.1-preview.2

Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model version. It follows the v1 branch as fixes land. To pin exact artifact bytes, pass a 40-character commit revision instead of 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] },
});