| { |
| "domain": "com.microsoft", |
| "name": "LinearAttentionGate", |
| "sinceVersion": 1, |
| "description": "Fuses the gate projections used by `com.microsoft.LinearAttention`'s gated-delta recurrence: `decay = decay_scale * softplus(a + dt_bias)` and, when requested, `beta = sigmoid(b)`. The last input axis is the head axis; `dt_bias` and `decay_scale` are float32 per-head vectors. Gate arithmetic is performed in float32 and narrowed only on store. Requesting `beta` requires `b`; an unconsumed `b` is permitted when `beta` is omitted.", |
| "inputs": [ |
| { |
| "role": "a", |
| "dtype": "T", |
| "description": "Decay gate projection with shape `(B, T, H)`. Any rank of at least 1 is accepted; the last axis is the head count and the leading axes are folded." |
| }, |
| { |
| "role": "dt_bias", |
| "dtype": "TF", |
| "rank": 1, |
| "description": "Per-head float32 bias added to `a`, with shape (H)." |
| }, |
| { |
| "role": "decay_scale", |
| "dtype": "TF", |
| "rank": 1, |
| "description": "Per-head float32 multiplier applied to `softplus(a + dt_bias)`, with shape `(H)`. For gated DeltaNet this is `-exp(A_log)`." |
| }, |
| { |
| "role": "b", |
| "dtype": "T", |
| "optional": true, |
| "description": "Update-rate projection with the same shape as `a` when `beta` is requested. It is accepted but unused when `beta` is omitted." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "decay", |
| "dtype": "T", |
| "rank": "ranks.aT", |
| "shape": "shapes.aT", |
| "description": "`decay_scale * softplus(a + dt_bias)`, with the same shape as `a`." |
| }, |
| { |
| "role": "beta", |
| "dtype": "T", |
| "rank": "ranks.aT", |
| "optional": true, |
| "shape": "shapes.aT", |
| "description": "sigmoid(b), with the same shape as `a`. Requires the `b` input." |
| } |
| ], |
| "typeConstraints": { "T": ["float32", "float16"], "TF": ["float32"] }, |
| "tunables": { "WORKGROUP_SIZE": 64 }, |
| "args": { |
| "aT": { "kind": "tensor", "semantic": "a", "role": "input" }, |
| "dtBiasT": { "kind": "tensor", "semantic": "dt_bias", "role": "weights" }, |
| "decayScaleT": { "kind": "tensor", "semantic": "decay_scale", "role": "weights" }, |
| "bT": { "kind": "tensor", "semantic": "b", "role": "input", "required": false }, |
| "decayT": { "kind": "tensor", "semantic": "decay", "role": "output" }, |
| "betaT": { "kind": "tensor", "semantic": "beta", "role": "output", "required": false } |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension", |
| "numHeads": "dim(shapes.aT, ranks.aT - 1)", |
| "gateCount": "numel(shapes.aT)", |
| "headsVec4": "numHeads / 4", |
| "gateVec4Count": "gateCount / 4", |
| "gateDtype": "tensorDtypes.aT", |
| "gateDtypeOk": "(gateDtype == \"float32\" or gateDtype == \"float16\") and f16Ok(dtypes.T)", |
| "paramsOk": "ranks.aT >= 1 and numHeads > 0 and ranks.dtBiasT == 1 and ranks.decayScaleT == 1 and tensorDtypes.dtBiasT == \"float32\" and tensorDtypes.decayScaleT == \"float32\" and dim(shapes.dtBiasT, 0) == numHeads and dim(shapes.decayScaleT, 0) == numHeads", |
| "tensorContract": "gateDtypeOk and paramsOk and sameShape(shapes.decayT, shapes.aT) and tensorDtypes.decayT == gateDtype", |
| "betaContract": "tensorContract and present.bT and present.betaT and sameShape(shapes.bT, shapes.aT) and sameShape(shapes.betaT, shapes.aT) and tensorDtypes.bT == gateDtype and tensorDtypes.betaT == gateDtype", |
| "decayOnlyContract": "tensorContract and not present.betaT", |
| "workgroupFits": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap", |
| "scalarDispatchFits": "ceilDiv(gateCount, tunables.WORKGROUP_SIZE) <= foldedDispatchCapacity", |
| "vec4DispatchFits": "numHeads % 4 == 0 and ceilDiv(gateVec4Count, tunables.WORKGROUP_SIZE) <= foldedDispatchCapacity" |
| }, |
| "bindingSets": { |
| "decayOnlyIo": [ |
| { |
| "name": "a", |
| "arg": "aT", |
| "semantic": "a", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$gateElement", |
| "length": "$gateItems" |
| }, |
| { |
| "name": "dt_bias", |
| "arg": "dtBiasT", |
| "semantic": "dt_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$paramElement", |
| "length": "$headItems" |
| }, |
| { |
| "name": "decay_scale", |
| "arg": "decayScaleT", |
| "semantic": "decay_scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$paramElement", |
| "length": "$headItems" |
| }, |
| { |
| "name": "decay", |
| "arg": "decayT", |
| "semantic": "decay", |
| "buffer": { "type": "storage" }, |
| "elementType": "$gateElement", |
| "length": "$gateItems" |
| } |
| ], |
| "withBetaIo": [ |
| { |
| "name": "a", |
| "arg": "aT", |
| "semantic": "a", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$gateElement", |
| "length": "$gateItems" |
| }, |
| { |
| "name": "dt_bias", |
| "arg": "dtBiasT", |
| "semantic": "dt_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$paramElement", |
| "length": "$headItems" |
| }, |
| { |
| "name": "decay_scale", |
| "arg": "decayScaleT", |
| "semantic": "decay_scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$paramElement", |
| "length": "$headItems" |
| }, |
| { |
| "name": "b", |
| "arg": "bT", |
| "semantic": "b", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$gateElement", |
| "length": "$gateItems" |
| }, |
| { |
| "name": "decay", |
| "arg": "decayT", |
| "semantic": "decay", |
| "buffer": { "type": "storage" }, |
| "elementType": "$gateElement", |
| "length": "$gateItems" |
| }, |
| { |
| "name": "beta", |
| "arg": "betaT", |
| "semantic": "beta", |
| "buffer": { "type": "storage" }, |
| "elementType": "$gateElement", |
| "length": "$gateItems" |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "vec4_gate_beta", |
| "priority": 30, |
| "when": ["betaContract", "workgroupFits", "vec4DispatchFits"], |
| "constants": { |
| "vectorized": true, |
| "hasBeta": true, |
| "usesF16": "gateDtype == \"float16\"", |
| "gateElement": "\"vec4<f16>\" if gateDtype == \"float16\" else \"vec4<f32>\"", |
| "paramElement": "\"vec4<f32>\"", |
| "headItems": "headsVec4", |
| "gateItems": "gateVec4Count", |
| "workgroupSize": "tunables.WORKGROUP_SIZE" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "LinearAttentionGate.Vec4GateBeta", |
| "shader": "linear-attention-gate.wgsl.jinja", |
| "bindings": "withBetaIo", |
| "dispatch": { "threads": "gateVec4Count", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| }, |
| { |
| "id": "vec4_gate", |
| "priority": 20, |
| "when": ["decayOnlyContract", "workgroupFits", "vec4DispatchFits"], |
| "constants": { |
| "vectorized": true, |
| "hasBeta": false, |
| "usesF16": "gateDtype == \"float16\"", |
| "gateElement": "\"vec4<f16>\" if gateDtype == \"float16\" else \"vec4<f32>\"", |
| "paramElement": "\"vec4<f32>\"", |
| "headItems": "headsVec4", |
| "gateItems": "gateVec4Count", |
| "workgroupSize": "tunables.WORKGROUP_SIZE" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "LinearAttentionGate.Vec4Gate", |
| "shader": "linear-attention-gate.wgsl.jinja", |
| "bindings": "decayOnlyIo", |
| "dispatch": { "threads": "gateVec4Count", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| }, |
| { |
| "id": "scalar_gate_beta", |
| "priority": 10, |
| "when": ["betaContract", "workgroupFits", "scalarDispatchFits"], |
| "constants": { |
| "vectorized": false, |
| "hasBeta": true, |
| "usesF16": "gateDtype == \"float16\"", |
| "gateElement": "\"f16\" if gateDtype == \"float16\" else \"f32\"", |
| "paramElement": "\"f32\"", |
| "headItems": "numHeads", |
| "gateItems": "gateCount", |
| "workgroupSize": "tunables.WORKGROUP_SIZE" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "LinearAttentionGate.ScalarGateBeta", |
| "shader": "linear-attention-gate.wgsl.jinja", |
| "bindings": "withBetaIo", |
| "dispatch": { "threads": "gateCount", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| }, |
| { |
| "id": "scalar_gate", |
| "priority": 0, |
| "when": ["decayOnlyContract", "workgroupFits", "scalarDispatchFits"], |
| "constants": { |
| "vectorized": false, |
| "hasBeta": false, |
| "usesF16": "gateDtype == \"float16\"", |
| "gateElement": "\"f16\" if gateDtype == \"float16\" else \"f32\"", |
| "paramElement": "\"f32\"", |
| "headItems": "numHeads", |
| "gateItems": "gateCount", |
| "workgroupSize": "tunables.WORKGROUP_SIZE" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "LinearAttentionGate.ScalarGate", |
| "shader": "linear-attention-gate.wgsl.jinja", |
| "bindings": "decayOnlyIo", |
| "dispatch": { "threads": "gateCount", "workgroupSize": "tunables.WORKGROUP_SIZE" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
