| { |
| "domain": "com.microsoft", |
| "name": "BiasSoftmax", |
| "sinceVersion": 1, |
| "description": "Computes `softmax(data + bias)` over the flattened suffix beginning at `axis`. The required `is_inner_broadcast` attribute selects how bias rows are reused: consecutive groups for inner broadcast or cyclic groups for outer broadcast. This specializes the `softmax(scores + additive_mask)` pattern used by transformer attention. Float16 and float32 are supported; the schema's double type is not.", |
| "inputs": [ |
| { "role": "data", "dtype": "T", "description": "The input data tensor." }, |
| { |
| "role": "bias", |
| "dtype": "T", |
| "description": "The bias (or additive mask) tensor. Its element count must be an integral number of flattened softmax rows and that row count must divide the data row count." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "rank": "ranks.data", |
| "shape": "shapes.data", |
| "description": "The output tensor; same shape as data." |
| } |
| ], |
| "attributes": { "axis": 1 }, |
| "attributeConstraints": { "is_inner_broadcast": { "required": true } }, |
| "attributeDescriptions": { |
| "axis": "The axis from which softmax is applied; dimensions from `axis` onward are included in the softmax reduction.", |
| "is_inner_broadcast": "When 1, bias is broadcast across dimensions from `broadcast_axis` to `axis-1`; when 0, bias is broadcast across dimensions 0 to `broadcast_axis-1`." |
| }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "args": { |
| "data": { "kind": "tensor", "semantic": "data", "role": "input" }, |
| "bias": { "kind": "tensor", "semantic": "bias", "role": "input" }, |
| "output": { "kind": "tensor", "semantic": "output", "role": "output" } |
| }, |
| "derive": { |
| "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data", |
| "batchCount": "outer(shapes.data, axisNorm)", |
| "blockSize": "dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm)", |
| "biasBlockCount": "numel(shapes.bias) / max(1, blockSize)", |
| "biasContract": "(numel(shapes.data) == 0 and numel(shapes.bias) == 0) or (blockSize > 0 and biasBlockCount > 0 and numel(shapes.bias) % blockSize == 0 and biasBlockCount <= batchCount and batchCount % biasBlockCount == 0)" |
| }, |
| "tunables": { "WORKGROUP_SIZE": 256, "BLOCK_COLS": 2048 }, |
| "bindingSets": { |
| "biasSoftmax": [ |
| { |
| "name": "data", |
| "arg": "data", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "bias", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "output", |
| "arg": "output", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "batchCount", "type": "u32", "value": "batchCount" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "longrow_split", |
| "priority": 40, |
| "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "blockSize >= 65536", "batchCount > 0", "batchCount <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(blockSize, tunables.BLOCK_COLS) <= device.limits.maxComputeWorkgroupsPerDimension", "f16Ok(dtypes.T)"], |
| "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "combineSubgroups": false }, |
| "intermediates": [ |
| { |
| "id": "blockMax", |
| "dtype": "float32", |
| "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]" |
| }, |
| { |
| "id": "blockSum", |
| "dtype": "float32", |
| "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]" |
| }, |
| { "id": "rowMax", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" }, |
| { "id": "rowSum", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" } |
| ], |
| "passes": [ |
| { |
| "id": "block_stats", |
| "name": "BiasSoftmax.LongRowBlockStats", |
| "source": { |
| "shader": "bias-softmax-longrow-stats.wgsl.jinja", |
| "inputs": { |
| "stage": "\"block\"", |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCount": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| } |
| }, |
| "bindings": [ |
| { |
| "name": "data", |
| "arg": "data", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "bias", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { "name": "blockMax", "semantic": "blockMax", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { "name": "blockSum", "semantic": "blockSum", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "blockSize", "type": "u32", "value": "blockSize" }, |
| { "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" } |
| ] |
| } |
| } |
| ], |
| "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" } |
| }, |
| { |
| "id": "row_stats", |
| "name": "BiasSoftmax.LongRowStats", |
| "source": { "shader": "bias-softmax-longrow-stats.wgsl.jinja", "inputs": { "stage": "\"row\"" } }, |
| "bindings": [ |
| { |
| "name": "blockMax", |
| "semantic": "blockMax", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "blockSum", |
| "semantic": "blockSum", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }] |
| } |
| } |
| ], |
| "dispatch": { "x": "batchCount" } |
| }, |
| { |
| "id": "normalize", |
| "name": "BiasSoftmax.LongRowNormalize", |
| "source": { |
| "shader": "bias-softmax-longrow-normalize.wgsl.jinja", |
| "inputs": { |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCount": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| } |
| }, |
| "bindings": [ |
| { |
| "name": "data", |
| "arg": "data", |
| "semantic": "data", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "bias", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$scalar" |
| }, |
| { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "output", |
| "arg": "output", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$scalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { "name": "Params", "fields": [{ "name": "blockSize", "type": "u32", "value": "blockSize" }] } |
| } |
| ], |
| "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rows", |
| "priority": 30, |
| "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "blockSize > 0", "blockSize <= 8", "batchCount >= 64", "f16Ok(dtypes.T)"], |
| "constants": { |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "combineSubgroups": false, |
| "packedRows": true |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.PackedRows", |
| "source": { |
| "shader": "bias-softmax.wgsl.jinja", |
| "inputs": { |
| "blockSize": "blockSize", |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCount": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| } |
| }, |
| "bindings": "biasSoftmax", |
| "dispatch": { "gridStride": "batchCount", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| }, |
| { |
| "id": "adaptive_row", |
| "priority": 10, |
| "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "numel(shapes.data) >= 0", "f16Ok(dtypes.T)"], |
| "constants": { |
| "packedRows": false, |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "combineSubgroups": "device.features.has(\"subgroups\")" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "BiasSoftmax.AdaptiveRow", |
| "source": { |
| "shader": "bias-softmax.wgsl.jinja", |
| "inputs": { |
| "blockSize": "blockSize", |
| "isInnerBroadcast": "attrs.is_inner_broadcast != 0", |
| "biasBlockCount": "max(1, biasBlockCount)", |
| "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))" |
| } |
| }, |
| "bindings": "biasSoftmax", |
| "dispatch": { "workgroups": "batchCount" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
