File size: 7,435 Bytes
81aba80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | {
"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)" }
}
]
}
]
}
|