| { |
| "domain": "com.microsoft", |
| "name": "GemmaRotaryEmbedding", |
| "sinceVersion": 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.", |
| "inputs": [ |
| { |
| "role": "emb", |
| "dtype": "U", |
| "rank": 3, |
| "description": "Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head." |
| }, |
| { |
| "role": "q", |
| "dtype": "T", |
| "rank": 4, |
| "description": "Query state with shape `(batch_size, num_heads, seq_len, dim)`." |
| }, |
| { "role": "q_rot", "dtype": "T", "rank": 4, "description": "Half-rotated query state, same shape as `q`." }, |
| { "role": "k", "dtype": "T", "rank": 4, "description": "Key state, same shape as `q`." }, |
| { "role": "k_rot", "dtype": "T", "rank": 4, "description": "Half-rotated key state, same shape as `q`." } |
| ], |
| "outputs": [ |
| { |
| "role": "output1", |
| "dtype": "T", |
| "rank": "ranks.qT", |
| "shape": "shapes.qT", |
| "description": "Rotary-embedded query, same shape as `q`." |
| }, |
| { |
| "role": "output2", |
| "dtype": "T", |
| "rank": "ranks.qT", |
| "shape": "shapes.qT", |
| "description": "Rotary-embedded key, same shape as `q`." |
| } |
| ], |
| "typeConstraints": { "T": ["float16"], "U": ["float32"] }, |
| "args": { |
| "embT": { "kind": "tensor", "semantic": "emb", "role": "input", "dtype": "float32" }, |
| "qT": { "kind": "tensor", "semantic": "q", "role": "input" }, |
| "qRotT": { "kind": "tensor", "semantic": "q_rot", "role": "input" }, |
| "kT": { "kind": "tensor", "semantic": "k", "role": "input" }, |
| "kRotT": { "kind": "tensor", "semantic": "k_rot", "role": "input" }, |
| "output1T": { "kind": "tensor", "semantic": "output1", "role": "output" }, |
| "output2T": { "kind": "tensor", "semantic": "output2", "role": "output" } |
| }, |
| "tunables": { "workgroupSize": 256 }, |
| "derive": { |
| "batchSize": "dim(shapes.qT, 0)", |
| "numHeads": "dim(shapes.qT, 1)", |
| "seqLen": "dim(shapes.qT, 2)", |
| "headDim": "dim(shapes.qT, 3)", |
| "sameAsQ": "ranks.qRotT == 4 and ranks.kT == 4 and ranks.kRotT == 4 and sameShape(shapes.qRotT, shapes.qT) and sameShape(shapes.kT, shapes.qT) and sameShape(shapes.kRotT, shapes.qT)", |
| "dtypesMatch": "tensorDtypes.qRotT == tensorDtypes.qT and tensorDtypes.kT == tensorDtypes.qT and tensorDtypes.kRotT == tensorDtypes.qT and tensorDtypes.output1T == tensorDtypes.qT and tensorDtypes.output2T == tensorDtypes.qT", |
| "contract": "ranks.qT == 4 and ranks.embT == 3 and sameAsQ and dtypesMatch and tensorDtypes.qT == \"float16\" and f16Ok(dtypes.T) and tensorDtypes.embT == \"float32\" and dim(shapes.embT, 0) == batchSize and dim(shapes.embT, 1) == seqLen and dim(shapes.embT, 2) == headDim and sameShape(shapes.output1T, shapes.qT) and sameShape(shapes.output2T, shapes.qT)", |
| "vec4Ok": "headDim % 4 == 0" |
| }, |
| "bindingSets": { |
| "scalar": [ |
| { |
| "name": "emb", |
| "arg": "embT", |
| "semantic": "emb", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "q", "arg": "qT", "semantic": "q", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, |
| { |
| "name": "q_rot", |
| "arg": "qRotT", |
| "semantic": "q_rot", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { "name": "k", "arg": "kT", "semantic": "k", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, |
| { |
| "name": "k_rot", |
| "arg": "kRotT", |
| "semantic": "k_rot", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "output1", |
| "arg": "output1T", |
| "semantic": "output1", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "output2", |
| "arg": "output2T", |
| "semantic": "output2", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "count", "type": "u32", "value": "numel(shapes.qT)" }, |
| { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim)" }, |
| { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim)" } |
| ] |
| } |
| } |
| ], |
| "vec4": [ |
| { |
| "name": "emb", |
| "arg": "embT", |
| "semantic": "emb", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { "name": "q", "arg": "qT", "semantic": "q", "buffer": { "type": "read-only-storage" }, "elementType": "$vector" }, |
| { |
| "name": "q_rot", |
| "arg": "qRotT", |
| "semantic": "q_rot", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$vector" |
| }, |
| { "name": "k", "arg": "kT", "semantic": "k", "buffer": { "type": "read-only-storage" }, "elementType": "$vector" }, |
| { |
| "name": "k_rot", |
| "arg": "kRotT", |
| "semantic": "k_rot", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$vector" |
| }, |
| { |
| "name": "output1", |
| "arg": "output1T", |
| "semantic": "output1", |
| "buffer": { "type": "storage" }, |
| "elementType": "$vector" |
| }, |
| { |
| "name": "output2", |
| "arg": "output2T", |
| "semantic": "output2", |
| "buffer": { "type": "storage" }, |
| "elementType": "$vector" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "count", "type": "u32", "value": "numel(shapes.qT) / 4" }, |
| { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim) / 4" }, |
| { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim) / 4" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "vec4", |
| "priority": 10, |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["contract", "vec4Ok", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"], |
| "constants": { |
| "vec4": true, |
| "scalar": "dtypes.T", |
| "vector": "\"vec4<f16>\"", |
| "workgroupSize": "tunables.workgroupSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmaRotaryEmbedding.Vec4", |
| "shader": "gemma-rotary-embedding.wgsl.jinja", |
| "bindings": "vec4", |
| "dispatch": { "threads": "numel(shapes.qT) / 4", "workgroupSize": "constants.workgroupSize" } |
| } |
| ] |
| }, |
| { |
| "id": "scalar", |
| "priority": 0, |
| "requires": { "features": ["shader-f16"] }, |
| "when": ["contract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"], |
| "constants": { |
| "vec4": false, |
| "scalar": "dtypes.T", |
| "vector": "\"vec4<f16>\"", |
| "workgroupSize": "tunables.workgroupSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GemmaRotaryEmbedding", |
| "shader": "gemma-rotary-embedding.wgsl.jinja", |
| "bindings": "scalar", |
| "dispatch": { "threads": "max(1, numel(shapes.qT))", "workgroupSize": "constants.workgroupSize" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
