Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.LinearAttention

Kernel card Files
xet
 Community

com.microsoft.LinearAttention

com.microsoft · ONNX Runtime contrib operator · contrib since_version 1

Description

Recurrent linear attention for packed [B, T, H*D] decode and prefill. It supports all four update rules, standard and inverse GQA, shared-key heads, and rollback states through state_window. Activations and state may independently use float16 or float32; bfloat16 is not implemented. past_state is optional for every update rule and defaults to zeros.

See the ONNX Runtime LinearAttention contrib-operator spec for the reference semantics.

Inputs

Name Bind key Logical dtype Rank Shape Description Presence
query queryT T 3 — Query vectors with 3D packed shape (B, T, H_q * d_k); heads are packed into the last dimension. required
key keyT T 3 — Key vectors with 3D packed shape (B, T, H_k * d_k), where positive H_k divides H_kv; H_k < H_kv shares each key head across multiple KV-state heads. Keys should be L2-normalized for delta/gated_delta modes. required
value valueT T 3 — Value vectors with 3D packed shape (B, T, H_kv * d_v). required
past_state pastStateT S derived derived; see description Recurrent state from the previous step with shape (B, H_kv, d_k, d_v), or (W, B, H_kv, d_k, d_v) when state_window = W > 0; defaults to zeros if absent. optional
decay decayT T 3 — Exponential decay gate in log-space with shape (B, T, H_kv * d_k) or (B, T, H_kv); required for gated and gated_delta modes. optional
beta betaT T 3 — Update rate (sigmoid output) with shape (B, T, H_kv) or (B, T, 1); required for delta and gated_delta modes. optional

Outputs

Name Bind key Logical dtype Rank Shape Description Presence
output outputT T 3 derived; see description Attention output with 3D packed shape (B, T, max(H_q, H_kv) * d_v). required
present_state presentStateT S derived derived; see description Updated recurrent state with shape (B, H_kv, d_k, d_v), or (W, B, H_kv, d_k, d_v) when state_window = W > 0. required

Attributes

Attributes and default values (overridable per request):

Attribute Default Description
chunk_size 64 Accepted for schema compatibility; does not affect the result.
scale 0 Scale applied to query-key products. Zero selects 1 / sqrt(d_k).
state_window 0 Number of recent recurrent states retained in present_state, in the supported range 0 to 8; zero returns only the current state.
update_rule "gated_delta" Recurrent update rule: linear, gated, delta, or gated_delta.
kv_num_heads — Number of key/value heads.
q_num_heads — Number of query heads.

Type constraints

Variable Allowed dtypes
T float32, float16
S float32, float16

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