| { |
| "domain": "com.microsoft", |
| "name": "GatedRelativePositionBias", |
| "sinceVersion": 1, |
| "inputs": { |
| "queryLayerT": { "onnx": "query_layer", "dtype": "T", "rank": 3 }, |
| "queryBiasT": { "onnx": "query_bias", "dtype": "T", "rank": 1 }, |
| "relPosT": { "onnx": "rel_pos", "dtype": "T", "rank": 4 }, |
| "weightT": { "onnx": "weight", "dtype": "T", "rank": 2 }, |
| "biasT": { "onnx": "bias", "dtype": "T", "rank": 1 }, |
| "ecoAT": { "onnx": "eco_a", "dtype": "T", "rank": 4 } |
| }, |
| "outputs": { |
| "outputT": { |
| "onnx": "output", |
| "dtype": "T", |
| "rank": 4, |
| "shape": "[dim(shapes.queryLayerT, 0), attrs.num_heads, dim(shapes.queryLayerT, 1), dim(shapes.queryLayerT, 1)]" |
| } |
| }, |
| "attributes": { "num_heads": {} }, |
| "attributeConstraints": { "num_heads": { "required": true } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { "WORKGROUP_SIZE": { "default": 128 }, "LANE_WORDS": { "default": 2 } }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32", |
| "batchSize": "dim(shapes.queryLayerT, 0)", |
| "seqLen": "dim(shapes.queryLayerT, 1)", |
| "queryHidden": "dim(shapes.queryLayerT, 2)", |
| "numHeads": "attrs.num_heads", |
| "headSize": "queryHidden / max(1, numHeads)", |
| "projDim": "dim(shapes.weightT, 1)", |
| "halfProj": "projDim / 2", |
| "rowCount": "batchSize * numHeads * seqLen", |
| "dtypeOk": "tensorDtypes.queryLayerT == tensorDtypes.queryBiasT and tensorDtypes.queryLayerT == tensorDtypes.relPosT and tensorDtypes.queryLayerT == tensorDtypes.weightT and tensorDtypes.queryLayerT == tensorDtypes.biasT and tensorDtypes.queryLayerT == tensorDtypes.ecoAT and tensorDtypes.queryLayerT == tensorDtypes.outputT and f16Ok(dtypes.T)", |
| "headsOk": "numHeads > 0 and queryHidden > 0 and queryHidden % numHeads == 0 and headSize > 0", |
| "projOk": "ranks.weightT == 2 and dim(shapes.weightT, 0) == headSize and projDim > 0 and projDim % 2 == 0 and projDim <= 32 and ranks.biasT == 1 and dim(shapes.biasT, 0) == projDim", |
| "relPosOk": "ranks.relPosT == 4 and dim(shapes.relPosT, 0) == 1 and dim(shapes.relPosT, 1) == numHeads and dim(shapes.relPosT, 2) == seqLen and dim(shapes.relPosT, 3) == seqLen", |
| "ecoOk": "ranks.ecoAT == 4 and numel(shapes.ecoAT) == numHeads", |
| "outputOk": "ranks.outputT == 4 and dim(shapes.outputT, 0) == batchSize and dim(shapes.outputT, 1) == numHeads and dim(shapes.outputT, 2) == seqLen and dim(shapes.outputT, 3) == seqLen", |
| "biasOk": "ranks.queryBiasT == 1 and dim(shapes.queryBiasT, 0) == queryHidden", |
| "contractOk": "ranks.queryLayerT == 3 and seqLen > 0 and dtypeOk and headsOk and projOk and relPosOk and ecoOk and outputOk and biasOk", |
| "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.LANE_WORDS > 0 and headSize + tunables.WORKGROUP_SIZE * 2 + 1 <= device.limits.maxComputeWorkgroupStorageSize / 4", |
| "dispatchOk": "rowCount <= foldedDispatchCapacity", |
| "vec4Ok": "seqLen % 4 == 0", |
| "vecWords": "seqLen / 4", |
| "vecWorkgroup": "max(32, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, ceilDiv(vecWords, max(1, tunables.LANE_WORDS))))))", |
| "scalarWorkgroup": "max(32, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, ceilDiv(seqLen, max(1, tunables.LANE_WORDS))))))", |
| "vecWordsPerLane": "ceilDiv(vecWords, max(1, vecWorkgroup))", |
| "scalarWordsPerLane": "ceilDiv(seqLen, max(1, scalarWorkgroup))", |
| "vectorElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "projectionLaneWidth": "min(vecWorkgroup, scalarWorkgroup, device.adapterInfo.subgroupMinSize) if (has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize > 0) else min(vecWorkgroup, scalarWorkgroup, 32)", |
| "headParallelPreferred": "min(headSize, projectionLaneWidth) >= 4 * min(projDim, projectionLaneWidth) or ((variableSubgroup16To32 or (not device.features.has(\"subgroups\") and not has(device.adapterInfo, \"subgroupMinSize\"))) and vec4Ok and headSize >= 8 * projDim and seqLen >= 2 * tunables.WORKGROUP_SIZE)", |
| "useVec4": "vec4Ok", |
| "rowWords": "vecWords if vec4Ok else seqLen", |
| "wordsPerLane": "vecWordsPerLane if vec4Ok else scalarWordsPerLane", |
| "stageRow": "wordsPerLane <= 8", |
| "workgroupSize": "vecWorkgroup if vec4Ok else scalarWorkgroup", |
| "scalar": "dtypes.T", |
| "vectorScalar": "vectorElement", |
| "rowElement": "vectorElement if vec4Ok else scalar" |
| }, |
| "when": ["contractOk", "workgroupOk", "dispatchOk"], |
| "bindings": { |
| "query_layer": { "arg": "queryLayerT", "elementType": "$scalar" }, |
| "query_bias": { "arg": "queryBiasT", "elementType": "$scalar" }, |
| "weight": { "arg": "weightT", "elementType": "$scalar" }, |
| "bias": { "arg": "biasT", "elementType": "$scalar" }, |
| "eco_a": { "arg": "ecoAT", "elementType": "$scalar" }, |
| "rel_pos": { "arg": "relPosT", "elementType": "$rowElement" }, |
| "output": { "arg": "outputT", "elementType": "$rowElement" }, |
| "params": { "struct": [{ "name": "seqLen", "type": "u32", "value": "seqLen" }] } |
| }, |
| "variants": [ |
| { |
| "id": "row", |
| "priority": 10, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRelativePositionBias.Row", |
| "shader": "gated-relative-position-bias.wgsl.jinja", |
| "bindings": ["query_layer", "query_bias", "weight", "bias", "eco_a", "rel_pos", "output", "params"], |
| "dispatch": { "x": "min(rowCount, DISPATCH_FOLD_WIDTH)", "y": "ceilDiv(rowCount, DISPATCH_FOLD_WIDTH)" } |
| } |
| ] |
| }, |
| { |
| "id": "row_head_tree", |
| "priority": 30, |
| "when": ["true"], |
| "demoteWhen": ["not headParallelPreferred"], |
| "requires": { "features": [] }, |
| "derive": { "parallelProjection": true, "useSubgroups": false }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRelativePositionBias.RowHeadTree", |
| "shader": "gated-relative-position-bias.wgsl.jinja", |
| "bindings": ["query_layer", "query_bias", "weight", "bias", "eco_a", "rel_pos", "output", "params"], |
| "dispatch": { "x": "min(rowCount, DISPATCH_FOLD_WIDTH)", "y": "ceilDiv(rowCount, DISPATCH_FOLD_WIDTH)" } |
| } |
| ] |
| }, |
| { |
| "id": "row_head_subgroups", |
| "priority": 40, |
| "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")"], |
| "demoteWhen": ["not headParallelPreferred"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { "parallelProjection": true, "useSubgroups": true }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRelativePositionBias.RowHeadSubgroups", |
| "shader": "gated-relative-position-bias.wgsl.jinja", |
| "bindings": ["query_layer", "query_bias", "weight", "bias", "eco_a", "rel_pos", "output", "params"], |
| "dispatch": { "x": "min(rowCount, DISPATCH_FOLD_WIDTH)", "y": "ceilDiv(rowCount, DISPATCH_FOLD_WIDTH)" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
