Xenova HF Staff

sync 2e7068faf55e

5b47652 verified 5 days ago

5.57 kB

library_name: kernels
license: apache-2.0
tags:
  - kernel
  - webgpu
  - wgsl

com.microsoft.EmbedLayerNormalization

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

Description

BERT embedding fusion: looks up word and position tables, optionally adds a segment table, then applies layer normalization. A segment table without IDs uses row 0. embedding_sum is the pre-normalization sum. mask_index is the first zero or the sequence length; without mask, it is zero. Batch and sequence dimensions must be non-empty.

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

Inputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`input_ids`	`inputIdsT`	`T1`	`int32`	`2`	—	Word ids of shape `(batch_size, sequence_length)`.	required
`segment_ids`	`segmentIdsT`	`T1`	`int32`	`2`	—	Segment ids `(batch_size, sequence_length)`. Requires `segment_embedding`; when omitted with that table present, every token uses row 0. Values must be valid non-negative table-row indices.	optional
`word_embedding`	`wordEmbeddingT`	`T`	same as logical dtype	`2`	—	Non-empty word embedding table `(vocab, hidden_size)`. Every `input_ids` value must be a valid non-negative row index.	required
`position_embedding`	`positionEmbeddingT`	`T`	same as logical dtype	`2`	—	Non-empty position embedding table `(max_positions, hidden_size)`. Without `position_ids`, it must contain at least `sequence_length` rows.	required
`segment_embedding`	`segmentEmbeddingT`	`T`	same as logical dtype	`2`	—	Non-empty segment embedding table `(segments, hidden_size)`. If `segment_ids` is absent, row 0 is used for every token.	optional
`gamma`	`gammaT`	`T`	same as logical dtype	`1`	—	Layer-normalization scale of shape `(hidden_size)`.	required
`beta`	`betaT`	`T`	same as logical dtype	`1`	—	Layer-normalization bias of shape `(hidden_size)`.	required
`mask`	`maskT`	`T1`	`int32`	`2`	—	Attention mask of shape `(batch_size, sequence_length)`. Only used to produce `mask_index`.	optional
`position_ids`	`positionIdsT`	`T1`	`int32`	`2`	—	Position ids `(batch_size, sequence_length)`, or `(1, sequence_length)` to share one row across the batch. Values must be valid non-negative table-row indices; absent uses the position within the sequence.	optional

Outputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`output`	`outputT`	`T`	same as logical dtype	`3`	derived; see description	Normalized embeddings of shape `(batch_size, sequence_length, hidden_size)`.	required
`mask_index`	`maskIndexT`	`T1`	`int32`	`1`	`[input_ids[0]]`	Position of the first zero in each mask row, or `sequence_length` when no zero exists; shape `(batch_size)`. It is zero when the optional mask input is absent.	optional
`embedding_sum`	`embeddingSumT`	`T`	same as logical dtype	`3`	derived; see description	The summed embeddings before normalization, including the segment term when present. Float16 uses staged `(word + segment) + position`; float32 uses `(word + position) + segment`.	optional

Attributes

Attributes and default values (overridable per request):

Attribute	Default	Description
`epsilon`	`9.999999960041972e-13`	Non-negative epsilon added to the layer-normalization variance before taking the square root.
`mask_index_type`	—	Optional shape-inference hint for the `mask_index` output type. The schema's `T1` constraint fixes the runtime tensor type to int32.

Type constraints

Variable	Allowed dtypes
`T`	`float32`, `float16`
`T1`	`int32`

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
embed-mask-index.wgsl.jinja
embed-normalize.wgsl.jinja
embed-sum.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/com.microsoft.EmbedLayerNormalization", { version: 1 });
const { outputT } = await kernel({
  inputIdsT: { data: inputIdsTData, shape: [1, 2] },
  wordEmbeddingT: { data: wordEmbeddingTData, shape: [2, 2] },
  positionEmbeddingT: { data: positionEmbeddingTData, shape: [2, 2] },
  gammaT: { data: gammaTData, shape: [2] },
  betaT: { data: betaTData, shape: [2] },
});