com.microsoft.GemmaRotaryEmbedding

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

Description

Fuses the Gemma rotary-embedding tail: computes sin and cos from float32 emb, casts them to float16, then evaluates q * cos + q_rot * sin and the corresponding expression for k. emb has shape (batch, seq, dim) and is broadcast over the head axis of the (batch, heads, seq, dim) operands. Each product is rounded to float16 before the addition.

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

Inputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`emb`	`embT`	`U`	`float32`	`3`	—	Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head.	required
`q`	`qT`	`T`	same as logical dtype	`4`	—	Query state with shape `(batch_size, num_heads, seq_len, dim)`.	required
`q_rot`	`qRotT`	`T`	same as logical dtype	`4`	—	Half-rotated query state, same shape as `q`.	required
`k`	`kT`	`T`	same as logical dtype	`4`	—	Key state, same shape as `q`.	required
`k_rot`	`kRotT`	`T`	same as logical dtype	`4`	—	Half-rotated key state, same shape as `q`.	required

Outputs

Name	Bind key	Logical dtype	Rank	Shape	Description	Presence
`output1`	`output1T`	`T`	same as `q`	same as `q`	Rotary-embedded query, same shape as `q`.	required
`output2`	`output2T`	`T`	same as `q`	same as `q`	Rotary-embedded key, same shape as `q`.	required

Type constraints

Variable	Allowed dtypes
`T`	`float16`
`U`	`float32`

Device requirements

Every implementation variant requires shader-f16; the package has no variant-level fallback without that capability.

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
gemma-rotary-embedding.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.GemmaRotaryEmbedding", { version: 1 });
const { output1T, output2T } = await kernel({
  embT: { data: embTData, shape: [1, 2, 4] },
  qT: { data: qTData, shape: [1, 1, 2, 4] },
  qRotT: { data: qRotTData, shape: [1, 1, 2, 4] },
  kT: { data: kTData, shape: [1, 1, 2, 4] },
  kRotT: { data: kRotTData, shape: [1, 1, 2, 4] },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.