| { |
| "domain": "com.microsoft", |
| "name": "EngramGate", |
| "sinceVersion": 1, |
| "inputs": { |
| "keyT": { "onnx": "key", "dtype": "T", "rank": 4 }, |
| "queryT": { "onnx": "query", "dtype": "T", "rank": 4 }, |
| "valueT": { "onnx": "value", "dtype": "T", "rank": 3 }, |
| "keyNormScaleT": { "onnx": "key_norm_scale", "dtype": "T", "rank": 2 }, |
| "queryNormScaleT": { "onnx": "query_norm_scale", "dtype": "T", "rank": 2 } |
| }, |
| "outputs": { "outputT": { "onnx": "output", "dtype": "T", "rank": 4, "shape": "shapes.keyT" } }, |
| "attributes": { "epsilon": { "default": 0.00001 } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "LANE_WORDS": { "default": 8 }, |
| "MIN_ROW_GROUPS": { "default": 8 } |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "hcMult": "dim(shapes.keyT, 2)", |
| "hidden": "dim(shapes.keyT, 3)", |
| "rowCount": "dim(shapes.keyT, 0) * dim(shapes.keyT, 1) * hcMult", |
| "dtypeOk": "tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.keyT == tensorDtypes.valueT and tensorDtypes.keyT == tensorDtypes.keyNormScaleT and tensorDtypes.keyT == tensorDtypes.queryNormScaleT and tensorDtypes.keyT == tensorDtypes.outputT and f16Ok(dtypes.T)", |
| "shapeOk": "ranks.keyT == 4 and ranks.valueT == 3 and hidden > 0 and hcMult > 0 and rowCount > 0 and sameShape(shapes.queryT, shapes.keyT) and sameShape(shapes.outputT, shapes.keyT) and sameShape(prefix(shapes.valueT, 2), prefix(shapes.keyT, 2)) and dim(shapes.valueT, 2) == hidden and sameShape(shapes.keyNormScaleT, suffix(shapes.keyT, 2)) and sameShape(shapes.queryNormScaleT, suffix(shapes.keyT, 2))", |
| "contractOk": "dtypeOk and shapeOk", |
| "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.LANE_WORDS > 0 and tunables.MIN_ROW_GROUPS > 0", |
| "sharedFoldFits": "tunables.WORKGROUP_SIZE * 16 <= device.limits.maxComputeWorkgroupStorageSize", |
| "vec4Ok": "hidden % 4 == 0", |
| "vecWords": "hidden / 4", |
| "vecPreferredLanes": "max(1, min(tunables.WORKGROUP_SIZE / 2, pow2ceil(max(1, ceilDiv(vecWords, max(1, tunables.LANE_WORDS))))))", |
| "vecWideLanes": "max(1, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, ceilDiv(vecWords, max(1, tunables.LANE_WORDS))))))", |
| "vecBaseLanes": "pick([[ceilDiv(vecWords, vecPreferredLanes) <= 8, vecPreferredLanes]], vecWideLanes)", |
| "vecBaseRowGroups": "ceilDiv(rowCount, max(1, tunables.WORKGROUP_SIZE / max(1, vecBaseLanes)))", |
| "vecFillLanes": "max(1, min(min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, vecWords))), pow2ceil(max(1, ceilDiv(tunables.MIN_ROW_GROUPS * tunables.WORKGROUP_SIZE, max(1, rowCount))))))", |
| "vecLanesPerRow": "pick([[vecBaseRowGroups < tunables.MIN_ROW_GROUPS, max(vecBaseLanes, vecFillLanes)]], vecBaseLanes)", |
| "vecRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, vecLanesPerRow)", |
| "vecWordsPerLane": "ceilDiv(vecWords, max(1, vecLanesPerRow))", |
| "vecRowGroups": "ceilDiv(rowCount, max(1, vecRowsPerGroup))", |
| "scalarPreferredLanes": "max(1, min(tunables.WORKGROUP_SIZE / 2, pow2ceil(max(1, ceilDiv(hidden, max(8, tunables.LANE_WORDS))))))", |
| "scalarWideLanes": "max(1, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, ceilDiv(hidden, max(1, tunables.LANE_WORDS))))))", |
| "scalarBaseLanes": "pick([[ceilDiv(hidden, scalarPreferredLanes) <= 8, scalarPreferredLanes]], scalarWideLanes)", |
| "scalarBaseRowGroups": "ceilDiv(rowCount, max(1, tunables.WORKGROUP_SIZE / max(1, scalarBaseLanes)))", |
| "scalarFillLanes": "max(1, min(min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, hidden))), pow2ceil(max(1, ceilDiv(tunables.MIN_ROW_GROUPS * tunables.WORKGROUP_SIZE, max(1, rowCount))))))", |
| "scalarLanesPerRow": "pick([[scalarBaseRowGroups < tunables.MIN_ROW_GROUPS, max(scalarBaseLanes, scalarFillLanes)]], scalarBaseLanes)", |
| "scalarRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, scalarLanesPerRow)", |
| "scalarWordsPerLane": "ceilDiv(hidden, max(1, scalarLanesPerRow))", |
| "scalarRowGroups": "ceilDiv(rowCount, max(1, scalarRowsPerGroup))", |
| "subgroupShapeKnown": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize > 0 and device.adapterInfo.subgroupMinSize <= device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMaxSize == 0", |
| "sgMin": "device.adapterInfo.subgroupMinSize if subgroupShapeKnown else 0", |
| "sgFullWidth": "subgroupShapeKnown and device.adapterInfo.subgroupMinSize == device.adapterInfo.subgroupMaxSize", |
| "sgVecLanesPerRow": "min(sgMin, vecLanesPerRow) if subgroupShapeKnown else 1", |
| "sgVecRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, sgVecLanesPerRow)", |
| "sgVecWordsPerLane": "ceilDiv(vecWords, max(1, sgVecLanesPerRow))", |
| "sgVecRowGroups": "ceilDiv(rowCount, max(1, sgVecRowsPerGroup))", |
| "sgScalarLanesPerRow": "min(sgMin, scalarLanesPerRow) if subgroupShapeKnown else 1", |
| "sgScalarRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, sgScalarLanesPerRow)", |
| "sgScalarWordsPerLane": "ceilDiv(hidden, max(1, sgScalarLanesPerRow))", |
| "sgScalarRowGroups": "ceilDiv(rowCount, max(1, sgScalarRowsPerGroup))", |
| "vectorElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "workgroupSize": "tunables.WORKGROUP_SIZE", |
| "epsilon": "attrs.epsilon", |
| "scalar": "dtypes.T", |
| "vectorScalar": "vectorElement" |
| }, |
| "when": ["contractOk", "workgroupOk"], |
| "bindings": { |
| "key": { "arg": "keyT", "elementType": "$vectorElement" }, |
| "query": { "arg": "queryT", "elementType": "$vectorElement" }, |
| "value": { "arg": "valueT", "elementType": "$vectorElement" }, |
| "kscale": { "arg": "keyNormScaleT", "elementType": "$vectorElement" }, |
| "qscale": { "arg": "queryNormScaleT", "elementType": "$vectorElement" }, |
| "out": { "arg": "outputT", "elementType": "$vectorElement" }, |
| "key_main": { "arg": "keyT", "name": "key", "elementType": "$scalar" }, |
| "query_main": { "arg": "queryT", "name": "query", "elementType": "$scalar" }, |
| "value_main": { "arg": "valueT", "name": "value", "elementType": "$scalar" }, |
| "kscale_main": { "arg": "keyNormScaleT", "name": "kscale", "elementType": "$scalar" }, |
| "qscale_main": { "arg": "queryNormScaleT", "name": "qscale", "elementType": "$scalar" }, |
| "out_main": { "arg": "outputT", "name": "out", "elementType": "$scalar" }, |
| "params": { "struct": [{ "name": "rows", "type": "u32", "value": "rowCount" }] } |
| }, |
| "variants": [ |
| { |
| "id": "subgroup_rows_vec4", |
| "priority": 30, |
| "when": ["vec4Ok", "subgroupShapeKnown", "sgVecWordsPerLane >= 1", "sgVecWordsPerLane <= 8", "sgVecRowGroups <= foldedDispatchCapacity"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "foldKind": "\"none\" if sgVecLanesPerRow <= 1 else (\"subgroup_full\" if (sgFullWidth and sgVecLanesPerRow == sgMin) else \"subgroup_butterfly\")", |
| "foldSteps": "log2ceil(sgVecLanesPerRow)", |
| "useVec4": "true", |
| "stageRow": "true", |
| "rowWords": "vecWords", |
| "wordsPerLane": "sgVecWordsPerLane", |
| "lanesPerRow": "sgVecLanesPerRow", |
| "rowsPerGroup": "sgVecRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "EngramGate.SubgroupRowsVec4", |
| "shader": "engram-gate-rows.wgsl.jinja", |
| "bindings": ["key", "query", "value", "kscale", "qscale", "out", "params"], |
| "dispatch": { |
| "x": "min(sgVecRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(sgVecRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rows_vec4", |
| "priority": 20, |
| "when": ["vec4Ok", "sharedFoldFits", "vecRowGroups <= foldedDispatchCapacity"], |
| "derive": { |
| "foldKind": "\"none\" if vecLanesPerRow <= 1 else \"workgroup_tree\"", |
| "foldSteps": "log2ceil(vecLanesPerRow)", |
| "useVec4": "true", |
| "stageRow": "vecWordsPerLane <= 8", |
| "rowWords": "vecWords", |
| "wordsPerLane": "vecWordsPerLane", |
| "lanesPerRow": "vecLanesPerRow", |
| "rowsPerGroup": "vecRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "EngramGate.PackedRowsVec4", |
| "shader": "engram-gate-rows.wgsl.jinja", |
| "bindings": ["key", "query", "value", "kscale", "qscale", "out", "params"], |
| "dispatch": { |
| "x": "min(vecRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(vecRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_rows_scalar", |
| "priority": 10, |
| "when": ["subgroupShapeKnown", "sgScalarWordsPerLane >= 1", "sgScalarWordsPerLane <= 8", "sgScalarRowGroups <= foldedDispatchCapacity"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "foldKind": "\"none\" if sgScalarLanesPerRow <= 1 else (\"subgroup_full\" if (sgFullWidth and sgScalarLanesPerRow == sgMin) else \"subgroup_butterfly\")", |
| "foldSteps": "log2ceil(sgScalarLanesPerRow)", |
| "useVec4": "false", |
| "stageRow": "true", |
| "rowWords": "hidden", |
| "wordsPerLane": "sgScalarWordsPerLane", |
| "lanesPerRow": "sgScalarLanesPerRow", |
| "rowsPerGroup": "sgScalarRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "EngramGate.SubgroupRowsScalar", |
| "shader": "engram-gate-rows.wgsl.jinja", |
| "bindings": ["key_main", "query_main", "value_main", "kscale_main", "qscale_main", "out_main", "params"], |
| "dispatch": { |
| "x": "min(sgScalarRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(sgScalarRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rows_scalar", |
| "priority": 0, |
| "when": ["sharedFoldFits", "scalarRowGroups <= foldedDispatchCapacity"], |
| "derive": { |
| "foldKind": "\"none\" if scalarLanesPerRow <= 1 else \"workgroup_tree\"", |
| "foldSteps": "log2ceil(scalarLanesPerRow)", |
| "useVec4": "false", |
| "stageRow": "scalarWordsPerLane <= 8", |
| "rowWords": "hidden", |
| "wordsPerLane": "scalarWordsPerLane", |
| "lanesPerRow": "scalarLanesPerRow", |
| "rowsPerGroup": "scalarRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "EngramGate.PackedRowsScalar", |
| "shader": "engram-gate-rows.wgsl.jinja", |
| "bindings": ["key_main", "query_main", "value_main", "kscale_main", "qscale_main", "out_main", "params"], |
| "dispatch": { |
| "x": "min(scalarRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(scalarRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
