Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.MultiHeadAttention

Kernel card Files
xet
 Community

com.microsoft.MultiHeadAttention

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

Description

Computes multi-head self- or cross-attention over explicit query, key, and value tensors, with an optional fused QKV projection bias and additive attention bias. Supports causal masking through unidirectional and a configurable score scale that defaults to 1 / sqrt(head_size). The current kernels require query, key, and value to have the same hidden width; distinct value-head widths, KV-cache, key-padding-mask, and diagnostic-QK modes are not yet implemented.

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

Inputs

Name Bind key Logical dtype Rank Shape Description Presence
query queryT T 3 — Query tensor of shape (batch_size, sequence_length, hidden_size). required
key keyT T 3 — Key tensor of shape (batch_size, kv_sequence_length, hidden_size); the current kernels require its hidden width to equal the query hidden width. required
value valueT T 3 — Value tensor of shape (batch_size, kv_sequence_length, v_hidden_size); the current supported subset requires v_hidden_size to equal the query hidden width. required
bias biasT T 1 — Optional fused bias from input projection with shape (hidden_size + hidden_size + v_hidden_size). optional
attention_bias attentionBiasT T 4 — Optional additive bias applied to the attention scores before softmax, of shape (batch_size or 1, num_heads or 1, sequence_length, kv_sequence_length); the first two dimensions broadcast. optional

Outputs

Name Bind key Logical dtype Rank Shape Description Presence
output outputT T 3 [query[0], query[1], value[2]] Attention output of shape (batch_size, sequence_length, v_hidden_size). required

Attributes

Attributes and default values (overridable per request):

Attribute Default Description
unidirectional 0 When 1, applies an upper-left causal mask: query position i attends only to key positions 0..min(i, kv_sequence_length - 1).
num_heads — Number of attention heads.
scale — Optional score scale; zero or omission selects 1 / sqrt(head_size).

Type constraints

Variable Allowed dtypes
T float32, float16

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/com.microsoft.MultiHeadAttention", { version: 1 });
const { outputT } = await kernel({
  queryT: { data: queryTData, shape: [2, 2, 4] },
  keyT: { data: keyTData, shape: [2, 3, 4] },
  valueT: { data: valueTData, shape: [2, 3, 4] },
}, {
  attrs: { num_heads: 2 },
});
Downloads last month
-
kernel
webgpu
wgsl
apache-2.0
WebGPU

Requires WebGPU support. See the compatibility table.