ai.onnx.LinearAttention

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

Description

Computes ONNX recurrent linear attention for packed [B, T, H*D] tensors. Standard GQA requires positive H_q divisible by H_kv; each KV head owns one state, and d_v must be positive. Supports all four update rules, optional rank-4 state, and float16 or float32 activations/state; bfloat16 is not implemented. Every positive d_k through 256 is supported; larger four-aligned widths are accepted when their required resources fit the WebGPU limits. Other head widths are unsupported.

See the ONNX LinearAttention spec for the reference semantics.

Inputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`query`	`queryT`	`T`	`3`	—	Query vectors with packed shape `(B, T, H_q * d_k)`; `H_q` must be a positive multiple of `H_kv`.	required
`key`	`keyT`	`T`	`3`	—	Key vectors with packed shape `(B, T, H_kv * d_k)`; keys should be L2-normalized for the `delta` update rules.	required
`value`	`valueT`	`T`	`3`	—	Value vectors with packed shape `(B, T, H_kv * d_v)`.	required
`past_state`	`pastStateT`	`S`	`4`	derived; see description	Optional recurrent state with shape `(B, H_kv, d_k, d_v)`. When absent, the recurrence starts from zero.	optional
`decay`	`decayT`	`T`	`3`	—	Log-space decay with shape `(B, T, H_kv * d_k)` or `(B, T, H_kv)`. Required only by `gated` and `gated_delta`.	optional
`beta`	`betaT`	`T`	`3`	—	Update rate with shape `(B, T, H_kv)` or `(B, T, 1)`. Required only by `delta` and `gated_delta`.	optional

Outputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`output`	`outputT`	`T`	`3`	derived; see description	Attention output with packed shape `(B, T, H_q * d_v)`.	required
`present_state`	`presentStateT`	`S`	`4`	derived; see description	Recurrent state after the final token, with shape `(B, H_kv, d_k, d_v)`. Its dtype matches `past_state` when supplied, otherwise it matches `query`.	required

Attributes

Attributes and default values (overridable per request):

Attribute	Default	Description
`chunk_size`	`64`	Correctness-neutral chunk-size hint for parallel prefill implementations. Every integer value has identical semantics.
`scale`	`0`	Output scaling factor. Zero selects `1 / sqrt(d_k)`.
`update_rule`	`"gated_delta"`	Recurrent update rule: `linear`, `gated`, `delta`, or `gated_delta`.
`kv_num_heads`	—	Required positive number of key/value heads.
`q_num_heads`	—	Required positive number of query heads; must be a multiple of `kv_num_heads`.

Type constraints

Variable	Allowed dtypes
`T`	`float32`, `float16`
`S`	`float32`, `float16`

Files

metadata.json — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
linear-attention.scalar.wgsl.jinja
linear-attention.vec4.wgsl.jinja

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.LinearAttention", { version: 1 });
const { outputT, presentStateT } = await kernel({
  queryT: { data: queryTData, shape: [1, 2, 8] },
  keyT: { data: keyTData, shape: [1, 2, 4] },
  valueT: { data: valueTData, shape: [1, 2, 3] },
  decayT: { data: decayTData, shape: [1, 2, 4] },
  betaT: { data: betaTData, shape: [1, 2, 1] },
}, {
  attrs: { q_num_heads: 2, kv_num_heads: 1 },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.