Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.MatMulNBitsMlp

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
3f9b2cb
·
verified ·
1 Parent(s): 8d4a681

sync 91d990483a17

Browse files
Files changed (7) hide show
  1. README.md +23 -19
  2. build/webgpu/bench.json +6 -6
  3. build/webgpu/manifest.json +229 -1506
  4. build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja +5 -10
  5. build/webgpu/metadata.json +39 -8
  6. build/webgpu/mlp-gate-up.wgsl.jinja +310 -47
  7. build/webgpu/test.json +349 -6
README.md CHANGED
@@ -18,24 +18,24 @@ See the [ONNX Runtime `MatMulNBitsMlp` contrib-operator spec](https://github.com
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
- | `A` | `aT` | `T1` | — | — | Shared activation of rank 2 `(M, K)` or rank 3 `(batch, sequence, K)`; only the last axis is the reduction axis. | required |
24
- | `skip` | `skipT` | `T1` | — | — | Residual added to `A` before normalization, with `A`'s shape. Requires `norm_scale`. | optional |
25
- | `norm_scale` | `normScaleT` | `T1` | `1` | — | Simplified-layer-normalization (RMS) gain of shape `[K]`. Absent means the projections read `A` unnormalized. | optional |
26
- | `gate_B` | `gateBT` | `uint8` | `3` | — | Bit-packed uint8 gate weights of shape `(N, k_blocks, blob_size)`. | required |
27
- | `gate_scales` | `gateScalesT` | `T1` | `2` | — | Per-block gate scales of shape `(N, k_blocks)`, with the same dtype as `A`. Quantization is symmetric: this operator has no zero-point input, so codes are offset by the midpoint `2^(bits - 1)`. | required |
28
- | `gate_bias` | `gateBiasT` | `T1` | `1` | — | Optional gate bias of shape `[N]`, added before the activation. | optional |
29
- | `up_B` | `upBT` | `uint8` | `3` | — | Bit-packed up weights, same shape and packing as gate_B. | required |
30
- | `up_scales` | `upScalesT` | `T1` | `2` | — | Per-block up scales of shape `(N, k_blocks)`. | required |
31
- | `up_bias` | `upBiasT` | `T1` | `1` | — | Optional up bias of shape `[N]`, added before the product. | optional |
32
 
33
  ## Outputs
34
 
35
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
36
  | --- | --- | --- | --- | --- | --- | --- |
37
- | `Y` | `yT` | `T1` | same as `A` | derived; see description | Gated MLP output: A's leading axes with a trailing N. | required |
38
- | `input_skip_bias_sum` | `residualT` | `T1` | same as `A` | same as `A` | The residual sum A + skip, with A's shape. Requires the skip input. | optional |
39
 
40
  ## Attributes
41
 
@@ -43,13 +43,13 @@ Attributes and default values (overridable per request):
43
 
44
  | Attribute | Default | Description |
45
  | --- | --- | --- |
46
- | `accuracy_level` | `0` | Minimum internal accuracy level: 0 (unset), 1 (float32), 2 (float16), 3 (bfloat16), or 4 (int8). |
47
- | `bits` | `4` | Bit width used to quantize both weight matrices; this implementation supports 2, 4, and 8. |
48
- | `epsilon` | `0.00001` | Epsilon used by the optional fused RMS normalization. |
49
  | `K` | — | Input feature dimension shared by both quantized weight matrices. |
50
  | `N` | — | Output feature dimension shared by both quantized weight matrices. |
 
51
  | `activation` | — | Activation applied to the gate projection; this implementation supports `silu`. |
 
52
  | `block_size` | — | Size of each quantization block along K. |
 
53
 
54
  ## Type constraints
55
 
@@ -59,7 +59,7 @@ Attributes and default values (overridable per request):
59
 
60
  ## Files
61
 
62
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
63
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
64
  - [`test.json`](build/webgpu/test.json) — correctness cases
65
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -68,10 +68,14 @@ Attributes and default values (overridable per request):
68
 
69
  ## Use with `@huggingface/kernels`
70
 
71
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
72
- It then allocates the result tensors automatically.
 
 
 
73
 
74
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
75
 
76
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
77
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
+ | `aT` | `A` | `T1` | — | — | Shared activation of rank 2 `(M, K)` or rank 3 `(batch, sequence, K)`; only the last axis is the reduction axis. | required |
24
+ | `skipT` | `skip` | `T1` | — | — | Residual added to `A` before normalization, with `A`'s shape. Requires `norm_scale`. | optional |
25
+ | `normScaleT` | `norm_scale` | `T1` | `1` | — | Simplified-layer-normalization (RMS) gain of shape `[K]`. Absent means the projections read `A` unnormalized. | optional |
26
+ | `gateBT` | `gate_B` | `uint8` | `3` | — | Bit-packed uint8 gate weights of shape `(N, k_blocks, blob_size)`. Bound in the packed storage layout: four blob bytes per u32 word. | required |
27
+ | `gateScalesT` | `gate_scales` | `T1` | `2` | — | Per-block gate scales of shape `(N, k_blocks)`, with the same dtype as `A`. Quantization is symmetric: this operator has no zero-point input, so codes are offset by the midpoint `2^(bits - 1)`. | required |
28
+ | `gateBiasT` | `gate_bias` | `T1` | `1` | — | Optional gate bias of shape `[N]`, added before the activation. | optional |
29
+ | `upBT` | `up_B` | `uint8` | `3` | — | Bit-packed up weights, same shape and packing as gate_B. Bound in the packed storage layout: four blob bytes per u32 word. | required |
30
+ | `upScalesT` | `up_scales` | `T1` | `2` | — | Per-block up scales of shape `(N, k_blocks)`. | required |
31
+ | `upBiasT` | `up_bias` | `T1` | `1` | — | Optional up bias of shape `[N]`, added before the product. | optional |
32
 
33
  ## Outputs
34
 
35
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
36
  | --- | --- | --- | --- | --- | --- | --- |
37
+ | `yT` | `Y` | `T1` | same as `aT` | derived | Gated MLP output: A's leading axes with a trailing N. | required |
38
+ | `residualT` | `input_skip_bias_sum` | `T1` | same as `aT` | same as `aT` | The residual sum A + skip, with A's shape. Requires the skip input. | optional |
39
 
40
  ## Attributes
41
 
 
43
 
44
  | Attribute | Default | Description |
45
  | --- | --- | --- |
 
 
 
46
  | `K` | — | Input feature dimension shared by both quantized weight matrices. |
47
  | `N` | — | Output feature dimension shared by both quantized weight matrices. |
48
+ | `accuracy_level` | `0` | Minimum internal accuracy level: 0 (unset), 1 (float32), 2 (float16), 3 (bfloat16), or 4 (int8). |
49
  | `activation` | — | Activation applied to the gate projection; this implementation supports `silu`. |
50
+ | `bits` | `4` | Bit width used to quantize both weight matrices; this implementation supports 2, 4, and 8. |
51
  | `block_size` | — | Size of each quantization block along K. |
52
+ | `epsilon` | `0.00001` | Epsilon used by the optional fused RMS normalization. |
53
 
54
  ## Type constraints
55
 
 
59
 
60
  ## Files
61
 
62
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
63
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
64
  - [`test.json`](build/webgpu/test.json) — correctness cases
65
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
68
 
69
  ## Use with `@huggingface/kernels`
70
 
71
+ ```sh
72
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
73
+ ```
74
+
75
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
76
 
77
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
78
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
79
 
80
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
81
 
build/webgpu/bench.json CHANGED
@@ -1,22 +1,22 @@
1
  {
2
- "op": "com.microsoft.MatMulNBitsMlp",
3
  "tunableSpace": { "TILE_N": [4, 8, 16], "LANES": [8, 16, 32] },
4
  "cases": [
5
  {
6
  "name": "mlp-q4-decode-k2048-n5632",
 
7
  "preset": "smoke",
8
  "vars": { "dtype": "float32" },
9
  "attrs": { "K": 2048, "N": 5632, "bits": 4, "block_size": 32, "activation": "silu" },
10
  "inputs": {
11
  "aT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 8101, "scale": 1 },
12
  "normScaleT": { "shape": [2048], "dtype": "float32", "dist": "normal", "seed": 8102, "scale": 1 },
13
- "gateBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8103, "scale": 255 },
14
  "gateScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8104, "scale": 0.05 },
15
- "upBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8105, "scale": 255 },
16
  "upScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8106, "scale": 0.05 }
17
  },
18
  "outputs": { "yT": { "shape": [1, 5632], "dtype": "float32" } },
19
- "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 5632 * 64 * 16 * 4" }] }
20
  },
21
  {
22
  "name": "mlp-q4-prefill-m64-k2048-n5632",
@@ -26,9 +26,9 @@
26
  "inputs": {
27
  "aT": { "shape": [64, 2048], "dtype": "float32", "dist": "normal", "seed": 8107, "scale": 1 },
28
  "normScaleT": { "shape": [2048], "dtype": "float32", "dist": "normal", "seed": 8108, "scale": 1 },
29
- "gateBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8109, "scale": 255 },
30
  "gateScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8110, "scale": 0.05 },
31
- "upBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8111, "scale": 255 },
32
  "upScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8112, "scale": 0.05 }
33
  },
34
  "outputs": { "yT": { "shape": [64, 5632], "dtype": "float32" } },
 
1
  {
 
2
  "tunableSpace": { "TILE_N": [4, 8, 16], "LANES": [8, 16, 32] },
3
  "cases": [
4
  {
5
  "name": "mlp-q4-decode-k2048-n5632",
6
+ "tunableSpace": {},
7
  "preset": "smoke",
8
  "vars": { "dtype": "float32" },
9
  "attrs": { "K": 2048, "N": 5632, "bits": 4, "block_size": 32, "activation": "silu" },
10
  "inputs": {
11
  "aT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 8101, "scale": 1 },
12
  "normScaleT": { "shape": [2048], "dtype": "float32", "dist": "normal", "seed": 8102, "scale": 1 },
13
+ "gateBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8103, "min": 0, "max": 256 },
14
  "gateScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8104, "scale": 0.05 },
15
+ "upBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8105, "min": 0, "max": 256 },
16
  "upScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8106, "scale": 0.05 }
17
  },
18
  "outputs": { "yT": { "shape": [1, 5632], "dtype": "float32" } },
19
+ "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 5632 * 64 * 16" }] }
20
  },
21
  {
22
  "name": "mlp-q4-prefill-m64-k2048-n5632",
 
26
  "inputs": {
27
  "aT": { "shape": [64, 2048], "dtype": "float32", "dist": "normal", "seed": 8107, "scale": 1 },
28
  "normScaleT": { "shape": [2048], "dtype": "float32", "dist": "normal", "seed": 8108, "scale": 1 },
29
+ "gateBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8109, "min": 0, "max": 256 },
30
  "gateScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8110, "scale": 0.05 },
31
+ "upBT": { "shape": [5632, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 8111, "min": 0, "max": 256 },
32
  "upScalesT": { "shape": [5632, 64], "dtype": "float32", "dist": "normal", "seed": 8112, "scale": 0.05 }
33
  },
34
  "outputs": { "yT": { "shape": [64, 5632], "dtype": "float32" } },
build/webgpu/manifest.json CHANGED
@@ -2,86 +2,35 @@
2
  "domain": "com.microsoft",
3
  "name": "MatMulNBitsMlp",
4
  "sinceVersion": 1,
5
- "description": "Fuses a gated MLP over two block-quantized projections that share one activation: `Y = silu(A_norm @ gate + gate_bias) * (A_norm @ up + up_bias)`, using the `MatMulNBits` weight packing with no zero-point input. `A_norm` is `A`, `SimplifiedLayerNormalization(A, norm_scale)`, or `SkipSimplifiedLayerNormalization(A, skip, norm_scale)`, whose residual sum may be returned as a second output. Only `silu` and the default `accuracy_level = 0` are implemented; bfloat16 is not implemented.",
6
- "inputs": [
7
- {
8
- "role": "A",
9
- "dtype": "T1",
10
- "description": "Shared activation of rank 2 `(M, K)` or rank 3 `(batch, sequence, K)`; only the last axis is the reduction axis."
11
- },
12
- {
13
- "role": "skip",
14
- "dtype": "T1",
15
- "optional": true,
16
- "description": "Residual added to `A` before normalization, with `A`'s shape. Requires `norm_scale`."
17
- },
18
- {
19
- "role": "norm_scale",
20
- "dtype": "T1",
21
- "rank": 1,
22
- "optional": true,
23
- "description": "Simplified-layer-normalization (RMS) gain of shape `[K]`. Absent means the projections read `A` unnormalized."
24
- },
25
- {
26
- "role": "gate_B",
27
- "dtype": "uint8",
28
- "rank": 3,
29
- "description": "Bit-packed uint8 gate weights of shape `(N, k_blocks, blob_size)`."
30
- },
31
- {
32
- "role": "gate_scales",
33
- "dtype": "T1",
34
- "rank": 2,
35
- "description": "Per-block gate scales of shape `(N, k_blocks)`, with the same dtype as `A`. Quantization is symmetric: this operator has no zero-point input, so codes are offset by the midpoint `2^(bits - 1)`."
36
- },
37
- {
38
- "role": "gate_bias",
39
- "dtype": "T1",
40
- "rank": 1,
41
- "optional": true,
42
- "description": "Optional gate bias of shape `[N]`, added before the activation."
43
- },
44
- {
45
- "role": "up_B",
46
- "dtype": "uint8",
47
- "rank": 3,
48
- "description": "Bit-packed up weights, same shape and packing as gate_B."
49
- },
50
- { "role": "up_scales", "dtype": "T1", "rank": 2, "description": "Per-block up scales of shape `(N, k_blocks)`." },
51
- {
52
- "role": "up_bias",
53
- "dtype": "T1",
54
- "rank": 1,
55
- "optional": true,
56
- "description": "Optional up bias of shape `[N]`, added before the product."
57
- }
58
- ],
59
- "outputs": [
60
- {
61
- "role": "Y",
62
- "dtype": "T1",
63
- "rank": "ranks.aT",
64
- "shape": "shapes.aT[:-1] + [attrs.N]",
65
- "description": "Gated MLP output: A's leading axes with a trailing N."
66
- },
67
- {
68
- "role": "input_skip_bias_sum",
69
  "dtype": "T1",
70
  "rank": "ranks.aT",
71
  "optional": true,
72
- "shape": "shapes.aT",
73
- "description": "The residual sum A + skip, with A's shape. Requires the skip input."
74
  }
75
- ],
76
- "attributes": { "accuracy_level": 0, "bits": 4, "epsilon": 0.00001 },
77
- "attributeDescriptions": {
78
- "K": "Input feature dimension shared by both quantized weight matrices.",
79
- "N": "Output feature dimension shared by both quantized weight matrices.",
80
- "accuracy_level": "Minimum internal accuracy level: 0 (unset), 1 (float32), 2 (float16), 3 (bfloat16), or 4 (int8).",
81
- "activation": "Activation applied to the gate projection; this implementation supports `silu`.",
82
- "bits": "Bit width used to quantize both weight matrices; this implementation supports 2, 4, and 8.",
83
- "block_size": "Size of each quantization block along K.",
84
- "epsilon": "Epsilon used by the optional fused RMS normalization."
85
  },
86
  "attributeConstraints": {
87
  "K": { "required": true },
@@ -92,24 +41,17 @@
92
  "block_size": { "required": true }
93
  },
94
  "typeConstraints": { "T1": ["float32", "float16"] },
95
- "args": {
96
- "aT": { "kind": "tensor", "semantic": "A", "role": "input" },
97
- "skipT": { "kind": "tensor", "semantic": "skip", "role": "input", "required": false },
98
- "normScaleT": { "kind": "tensor", "semantic": "norm_scale", "role": "weights", "required": false },
99
- "gateBT": { "kind": "tensor", "semantic": "gate_B", "role": "weights" },
100
- "gateScalesT": { "kind": "tensor", "semantic": "gate_scales", "role": "weights" },
101
- "gateBiasT": { "kind": "tensor", "semantic": "gate_bias", "role": "weights", "required": false },
102
- "upBT": { "kind": "tensor", "semantic": "up_B", "role": "weights" },
103
- "upScalesT": { "kind": "tensor", "semantic": "up_scales", "role": "weights" },
104
- "upBiasT": { "kind": "tensor", "semantic": "up_bias", "role": "weights", "required": false },
105
- "yT": { "kind": "tensor", "semantic": "Y", "role": "output" },
106
- "residualT": { "kind": "tensor", "semantic": "input_skip_bias_sum", "role": "output", "required": false }
107
  },
108
- "tunables": { "TILE_N": 8, "LANES": 8, "NORM_WORKGROUP_SIZE": 128, "ROW_TILE": 8 },
109
  "derive": {
110
  "aRows": "numel(shapes.aT) / max(1, attrs.K)",
111
- "rowTile": "1 if aRows <= 1 else min(aRows, tunables.ROW_TILE)",
112
- "rowGroups": "ceilDiv(aRows, rowTile)",
113
  "kBlocks": "dim(shapes.gateBT, 1)",
114
  "blobSize": "dim(shapes.gateBT, 2)",
115
  "codesPerByte": "8 / attrs.bits",
@@ -124,30 +66,38 @@
124
  "lanesPow2": "tunables.LANES == pow2ceil(tunables.LANES)",
125
  "mlpShapeOk": "bitsSupported and weightShapeOk and scaleShapeOk and ioShapeOk and biasShapeOk and dtypeOk and lanesPow2 and attrs.K > 0 and attrs.N > 0 and attrs.block_size > 0",
126
  "normContractOk": "present.normScaleT and ranks.normScaleT == 1 and dim(shapes.normScaleT, 0) == attrs.K and tensorDtypes.normScaleT == tensorDtypes.aT and (sameShape(shapes.skipT, shapes.aT) and tensorDtypes.skipT == tensorDtypes.aT if present.skipT else true) and (sameShape(shapes.residualT, shapes.aT) and tensorDtypes.residualT == tensorDtypes.aT and present.skipT if present.residualT else true)",
127
- "gateUpDispatchFits": "ceilDiv(attrs.N, tunables.TILE_N) <= device.limits.maxComputeWorkgroupsPerDimension and aRows <= device.limits.maxComputeWorkgroupsPerDimension and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX",
 
 
 
 
 
 
 
 
128
  "normDispatchFits": "tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX",
129
  "biasPresence_nogb_noub": "not present.gateBiasT and not present.upBiasT",
130
  "biasPresence_nogb_ub": "not present.gateBiasT and present.upBiasT",
131
  "biasPresence_gb_noub": "present.gateBiasT and not present.upBiasT",
132
- "biasPresence_gb_ub": "present.gateBiasT and present.upBiasT"
133
- },
134
- "constants": {
135
  "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
136
  "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
137
- "usesF16": "tensorDtypes.aT == \"float16\"",
138
  "K": "attrs.K",
139
  "N": "attrs.N",
140
  "blockSize": "attrs.block_size",
141
- "kBlocks": "kBlocks",
142
- "blobSize": "blobSize",
143
  "bits": "attrs.bits",
144
- "codesPerByte": "codesPerByte",
145
- "codeMask": "codeMask",
146
  "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
147
  "tileN": "tunables.TILE_N",
148
  "lanes": "tunables.LANES",
149
- "rowTile": "rowTile",
150
- "rows": "aRows",
 
 
 
 
 
 
 
151
  "hidden": "attrs.K",
152
  "workgroupSize": "tunables.NORM_WORKGROUP_SIZE",
153
  "epsilon": "epsilonValue",
@@ -158,1789 +108,562 @@
158
  "K_LEN": "attrs.K",
159
  "N_LEN": "attrs.N"
160
  },
161
- "bindingSets": {
162
- "normFull": [
163
- {
164
- "name": "a",
165
- "arg": "aT",
166
- "semantic": "A",
167
- "buffer": { "type": "read-only-storage" },
168
- "elementType": "$aScalar"
169
- },
170
- {
171
- "name": "skip",
172
- "arg": "skipT",
173
- "semantic": "skip",
174
- "buffer": { "type": "read-only-storage" },
175
- "elementType": "$aScalar"
176
- },
177
- {
178
- "name": "norm_scale",
179
- "arg": "normScaleT",
180
- "semantic": "norm_scale",
181
- "buffer": { "type": "read-only-storage" },
182
- "elementType": "$aScalar",
183
- "length": "$K_LEN"
184
- },
185
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "storage" }, "elementType": "f32" },
186
- {
187
- "name": "residual",
188
- "arg": "residualT",
189
- "semantic": "input_skip_bias_sum",
190
- "buffer": { "type": "storage" },
191
- "elementType": "$aScalar"
192
- },
193
- {
194
- "name": "params",
195
- "semantic": "kernel.params",
196
- "buffer": { "type": "uniform" },
197
- "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "aRows" }] }
198
- }
199
- ],
200
- "normSkip": [
201
- {
202
- "name": "a",
203
- "arg": "aT",
204
- "semantic": "A",
205
- "buffer": { "type": "read-only-storage" },
206
- "elementType": "$aScalar"
207
- },
208
- {
209
- "name": "skip",
210
- "arg": "skipT",
211
- "semantic": "skip",
212
- "buffer": { "type": "read-only-storage" },
213
- "elementType": "$aScalar"
214
- },
215
- {
216
- "name": "norm_scale",
217
- "arg": "normScaleT",
218
- "semantic": "norm_scale",
219
- "buffer": { "type": "read-only-storage" },
220
- "elementType": "$aScalar",
221
- "length": "$K_LEN"
222
- },
223
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "storage" }, "elementType": "f32" },
224
- {
225
- "name": "params",
226
- "semantic": "kernel.params",
227
- "buffer": { "type": "uniform" },
228
- "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "aRows" }] }
229
- }
230
- ],
231
- "normOnly": [
232
- {
233
- "name": "a",
234
- "arg": "aT",
235
- "semantic": "A",
236
- "buffer": { "type": "read-only-storage" },
237
- "elementType": "$aScalar"
238
- },
239
- {
240
- "name": "norm_scale",
241
- "arg": "normScaleT",
242
- "semantic": "norm_scale",
243
- "buffer": { "type": "read-only-storage" },
244
- "elementType": "$aScalar",
245
- "length": "$K_LEN"
246
- },
247
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "storage" }, "elementType": "f32" },
248
- {
249
- "name": "params",
250
- "semantic": "kernel.params",
251
- "buffer": { "type": "uniform" },
252
- "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "aRows" }] }
253
- }
254
- ],
255
- "gateUpAFull": [
256
- {
257
- "name": "a",
258
- "arg": "aT",
259
- "semantic": "A",
260
- "buffer": { "type": "read-only-storage" },
261
- "elementType": "$aScalar"
262
- },
263
- {
264
- "name": "gate_b",
265
- "arg": "gateBT",
266
- "semantic": "gate_B",
267
- "buffer": { "type": "read-only-storage" },
268
- "elementType": "u32"
269
- },
270
- {
271
- "name": "gate_scales",
272
- "arg": "gateScalesT",
273
- "semantic": "gate_scales",
274
- "buffer": { "type": "read-only-storage" },
275
- "elementType": "$aScalar"
276
- },
277
- {
278
- "name": "gate_bias",
279
- "arg": "gateBiasT",
280
- "semantic": "gate_bias",
281
- "buffer": { "type": "read-only-storage" },
282
- "elementType": "$aScalar",
283
- "length": "$N_LEN"
284
- },
285
- {
286
- "name": "up_b",
287
- "arg": "upBT",
288
- "semantic": "up_B",
289
- "buffer": { "type": "read-only-storage" },
290
- "elementType": "u32"
291
- },
292
- {
293
- "name": "up_scales",
294
- "arg": "upScalesT",
295
- "semantic": "up_scales",
296
- "buffer": { "type": "read-only-storage" },
297
- "elementType": "$aScalar"
298
- },
299
- {
300
- "name": "up_bias",
301
- "arg": "upBiasT",
302
- "semantic": "up_bias",
303
- "buffer": { "type": "read-only-storage" },
304
- "elementType": "$aScalar",
305
- "length": "$N_LEN"
306
- },
307
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
308
- ],
309
- "gateUpNormedFull": [
310
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
311
- {
312
- "name": "gate_b",
313
- "arg": "gateBT",
314
- "semantic": "gate_B",
315
- "buffer": { "type": "read-only-storage" },
316
- "elementType": "u32"
317
- },
318
- {
319
- "name": "gate_scales",
320
- "arg": "gateScalesT",
321
- "semantic": "gate_scales",
322
- "buffer": { "type": "read-only-storage" },
323
- "elementType": "$aScalar"
324
- },
325
- {
326
- "name": "gate_bias",
327
- "arg": "gateBiasT",
328
- "semantic": "gate_bias",
329
- "buffer": { "type": "read-only-storage" },
330
- "elementType": "$aScalar",
331
- "length": "$N_LEN"
332
- },
333
- {
334
- "name": "up_b",
335
- "arg": "upBT",
336
- "semantic": "up_B",
337
- "buffer": { "type": "read-only-storage" },
338
- "elementType": "u32"
339
- },
340
- {
341
- "name": "up_scales",
342
- "arg": "upScalesT",
343
- "semantic": "up_scales",
344
- "buffer": { "type": "read-only-storage" },
345
- "elementType": "$aScalar"
346
- },
347
- {
348
- "name": "up_bias",
349
- "arg": "upBiasT",
350
- "semantic": "up_bias",
351
- "buffer": { "type": "read-only-storage" },
352
- "elementType": "$aScalar",
353
- "length": "$N_LEN"
354
- },
355
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
356
- ],
357
- "fusedFull": [
358
- {
359
- "name": "a",
360
- "arg": "aT",
361
- "semantic": "A",
362
- "buffer": { "type": "read-only-storage" },
363
- "elementType": "$aScalar"
364
- },
365
- {
366
- "name": "skip",
367
- "arg": "skipT",
368
- "semantic": "skip",
369
- "buffer": { "type": "read-only-storage" },
370
- "elementType": "$aScalar"
371
- },
372
- {
373
- "name": "norm_scale",
374
- "arg": "normScaleT",
375
- "semantic": "norm_scale",
376
- "buffer": { "type": "read-only-storage" },
377
- "elementType": "$aScalar",
378
- "length": "$K_LEN"
379
- },
380
- {
381
- "name": "gate_b",
382
- "arg": "gateBT",
383
- "semantic": "gate_B",
384
- "buffer": { "type": "read-only-storage" },
385
- "elementType": "u32"
386
- },
387
- {
388
- "name": "gate_scales",
389
- "arg": "gateScalesT",
390
- "semantic": "gate_scales",
391
- "buffer": { "type": "read-only-storage" },
392
- "elementType": "$aScalar"
393
- },
394
- {
395
- "name": "gate_bias",
396
- "arg": "gateBiasT",
397
- "semantic": "gate_bias",
398
- "buffer": { "type": "read-only-storage" },
399
- "elementType": "$aScalar",
400
- "length": "$N_LEN"
401
- },
402
- {
403
- "name": "up_b",
404
- "arg": "upBT",
405
- "semantic": "up_B",
406
- "buffer": { "type": "read-only-storage" },
407
- "elementType": "u32"
408
- },
409
- {
410
- "name": "up_scales",
411
- "arg": "upScalesT",
412
- "semantic": "up_scales",
413
- "buffer": { "type": "read-only-storage" },
414
- "elementType": "$aScalar"
415
- },
416
- {
417
- "name": "up_bias",
418
- "arg": "upBiasT",
419
- "semantic": "up_bias",
420
- "buffer": { "type": "read-only-storage" },
421
- "elementType": "$aScalar",
422
- "length": "$N_LEN"
423
- },
424
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
425
- {
426
- "name": "residual",
427
- "arg": "residualT",
428
- "semantic": "input_skip_bias_sum",
429
- "buffer": { "type": "storage" },
430
- "elementType": "$aScalar"
431
- }
432
- ],
433
- "gateUpA_nogb_noub": [
434
- {
435
- "name": "a",
436
- "arg": "aT",
437
- "semantic": "A",
438
- "buffer": { "type": "read-only-storage" },
439
- "elementType": "$aScalar"
440
- },
441
- {
442
- "name": "gate_b",
443
- "arg": "gateBT",
444
- "semantic": "gate_B",
445
- "buffer": { "type": "read-only-storage" },
446
- "elementType": "u32"
447
- },
448
- {
449
- "name": "gate_scales",
450
- "arg": "gateScalesT",
451
- "semantic": "gate_scales",
452
- "buffer": { "type": "read-only-storage" },
453
- "elementType": "$aScalar"
454
- },
455
- {
456
- "name": "up_b",
457
- "arg": "upBT",
458
- "semantic": "up_B",
459
- "buffer": { "type": "read-only-storage" },
460
- "elementType": "u32"
461
- },
462
- {
463
- "name": "up_scales",
464
- "arg": "upScalesT",
465
- "semantic": "up_scales",
466
- "buffer": { "type": "read-only-storage" },
467
- "elementType": "$aScalar"
468
- },
469
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
470
- ],
471
- "gateUpNormed_nogb_noub": [
472
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
473
- {
474
- "name": "gate_b",
475
- "arg": "gateBT",
476
- "semantic": "gate_B",
477
- "buffer": { "type": "read-only-storage" },
478
- "elementType": "u32"
479
- },
480
- {
481
- "name": "gate_scales",
482
- "arg": "gateScalesT",
483
- "semantic": "gate_scales",
484
- "buffer": { "type": "read-only-storage" },
485
- "elementType": "$aScalar"
486
- },
487
- {
488
- "name": "up_b",
489
- "arg": "upBT",
490
- "semantic": "up_B",
491
- "buffer": { "type": "read-only-storage" },
492
- "elementType": "u32"
493
- },
494
- {
495
- "name": "up_scales",
496
- "arg": "upScalesT",
497
- "semantic": "up_scales",
498
- "buffer": { "type": "read-only-storage" },
499
- "elementType": "$aScalar"
500
- },
501
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
502
- ],
503
- "fused_norm_nogb_noub": [
504
- {
505
- "name": "a",
506
- "arg": "aT",
507
- "semantic": "A",
508
- "buffer": { "type": "read-only-storage" },
509
- "elementType": "$aScalar"
510
- },
511
- {
512
- "name": "norm_scale",
513
- "arg": "normScaleT",
514
- "semantic": "norm_scale",
515
- "buffer": { "type": "read-only-storage" },
516
- "elementType": "$aScalar",
517
- "length": "$K_LEN"
518
- },
519
- {
520
- "name": "gate_b",
521
- "arg": "gateBT",
522
- "semantic": "gate_B",
523
- "buffer": { "type": "read-only-storage" },
524
- "elementType": "u32"
525
- },
526
- {
527
- "name": "gate_scales",
528
- "arg": "gateScalesT",
529
- "semantic": "gate_scales",
530
- "buffer": { "type": "read-only-storage" },
531
- "elementType": "$aScalar"
532
- },
533
- {
534
- "name": "up_b",
535
- "arg": "upBT",
536
- "semantic": "up_B",
537
- "buffer": { "type": "read-only-storage" },
538
- "elementType": "u32"
539
- },
540
- {
541
- "name": "up_scales",
542
- "arg": "upScalesT",
543
- "semantic": "up_scales",
544
- "buffer": { "type": "read-only-storage" },
545
- "elementType": "$aScalar"
546
- },
547
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
548
- ],
549
- "fused_skip_nogb_noub": [
550
- {
551
- "name": "a",
552
- "arg": "aT",
553
- "semantic": "A",
554
- "buffer": { "type": "read-only-storage" },
555
- "elementType": "$aScalar"
556
- },
557
- {
558
- "name": "skip",
559
- "arg": "skipT",
560
- "semantic": "skip",
561
- "buffer": { "type": "read-only-storage" },
562
- "elementType": "$aScalar"
563
- },
564
- {
565
- "name": "norm_scale",
566
- "arg": "normScaleT",
567
- "semantic": "norm_scale",
568
- "buffer": { "type": "read-only-storage" },
569
- "elementType": "$aScalar",
570
- "length": "$K_LEN"
571
- },
572
- {
573
- "name": "gate_b",
574
- "arg": "gateBT",
575
- "semantic": "gate_B",
576
- "buffer": { "type": "read-only-storage" },
577
- "elementType": "u32"
578
- },
579
- {
580
- "name": "gate_scales",
581
- "arg": "gateScalesT",
582
- "semantic": "gate_scales",
583
- "buffer": { "type": "read-only-storage" },
584
- "elementType": "$aScalar"
585
- },
586
- {
587
- "name": "up_b",
588
- "arg": "upBT",
589
- "semantic": "up_B",
590
- "buffer": { "type": "read-only-storage" },
591
- "elementType": "u32"
592
- },
593
- {
594
- "name": "up_scales",
595
- "arg": "upScalesT",
596
- "semantic": "up_scales",
597
- "buffer": { "type": "read-only-storage" },
598
- "elementType": "$aScalar"
599
- },
600
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
601
- ],
602
- "fused_skipsum_nogb_noub": [
603
- {
604
- "name": "a",
605
- "arg": "aT",
606
- "semantic": "A",
607
- "buffer": { "type": "read-only-storage" },
608
- "elementType": "$aScalar"
609
- },
610
- {
611
- "name": "skip",
612
- "arg": "skipT",
613
- "semantic": "skip",
614
- "buffer": { "type": "read-only-storage" },
615
- "elementType": "$aScalar"
616
- },
617
- {
618
- "name": "norm_scale",
619
- "arg": "normScaleT",
620
- "semantic": "norm_scale",
621
- "buffer": { "type": "read-only-storage" },
622
- "elementType": "$aScalar",
623
- "length": "$K_LEN"
624
- },
625
- {
626
- "name": "gate_b",
627
- "arg": "gateBT",
628
- "semantic": "gate_B",
629
- "buffer": { "type": "read-only-storage" },
630
- "elementType": "u32"
631
- },
632
- {
633
- "name": "gate_scales",
634
- "arg": "gateScalesT",
635
- "semantic": "gate_scales",
636
- "buffer": { "type": "read-only-storage" },
637
- "elementType": "$aScalar"
638
- },
639
- {
640
- "name": "up_b",
641
- "arg": "upBT",
642
- "semantic": "up_B",
643
- "buffer": { "type": "read-only-storage" },
644
- "elementType": "u32"
645
- },
646
- {
647
- "name": "up_scales",
648
- "arg": "upScalesT",
649
- "semantic": "up_scales",
650
- "buffer": { "type": "read-only-storage" },
651
- "elementType": "$aScalar"
652
- },
653
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
654
- {
655
- "name": "residual",
656
- "arg": "residualT",
657
- "semantic": "input_skip_bias_sum",
658
- "buffer": { "type": "storage" },
659
- "elementType": "$aScalar"
660
- }
661
- ],
662
- "gateUpA_nogb_ub": [
663
- {
664
- "name": "a",
665
- "arg": "aT",
666
- "semantic": "A",
667
- "buffer": { "type": "read-only-storage" },
668
- "elementType": "$aScalar"
669
- },
670
- {
671
- "name": "gate_b",
672
- "arg": "gateBT",
673
- "semantic": "gate_B",
674
- "buffer": { "type": "read-only-storage" },
675
- "elementType": "u32"
676
- },
677
- {
678
- "name": "gate_scales",
679
- "arg": "gateScalesT",
680
- "semantic": "gate_scales",
681
- "buffer": { "type": "read-only-storage" },
682
- "elementType": "$aScalar"
683
- },
684
- {
685
- "name": "up_b",
686
- "arg": "upBT",
687
- "semantic": "up_B",
688
- "buffer": { "type": "read-only-storage" },
689
- "elementType": "u32"
690
- },
691
- {
692
- "name": "up_scales",
693
- "arg": "upScalesT",
694
- "semantic": "up_scales",
695
- "buffer": { "type": "read-only-storage" },
696
- "elementType": "$aScalar"
697
- },
698
- {
699
- "name": "up_bias",
700
- "arg": "upBiasT",
701
- "semantic": "up_bias",
702
- "buffer": { "type": "read-only-storage" },
703
- "elementType": "$aScalar",
704
- "length": "$N_LEN"
705
- },
706
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
707
- ],
708
- "gateUpNormed_nogb_ub": [
709
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
710
- {
711
- "name": "gate_b",
712
- "arg": "gateBT",
713
- "semantic": "gate_B",
714
- "buffer": { "type": "read-only-storage" },
715
- "elementType": "u32"
716
- },
717
- {
718
- "name": "gate_scales",
719
- "arg": "gateScalesT",
720
- "semantic": "gate_scales",
721
- "buffer": { "type": "read-only-storage" },
722
- "elementType": "$aScalar"
723
- },
724
- {
725
- "name": "up_b",
726
- "arg": "upBT",
727
- "semantic": "up_B",
728
- "buffer": { "type": "read-only-storage" },
729
- "elementType": "u32"
730
- },
731
- {
732
- "name": "up_scales",
733
- "arg": "upScalesT",
734
- "semantic": "up_scales",
735
- "buffer": { "type": "read-only-storage" },
736
- "elementType": "$aScalar"
737
- },
738
- {
739
- "name": "up_bias",
740
- "arg": "upBiasT",
741
- "semantic": "up_bias",
742
- "buffer": { "type": "read-only-storage" },
743
- "elementType": "$aScalar",
744
- "length": "$N_LEN"
745
- },
746
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
747
- ],
748
- "fused_norm_nogb_ub": [
749
- {
750
- "name": "a",
751
- "arg": "aT",
752
- "semantic": "A",
753
- "buffer": { "type": "read-only-storage" },
754
- "elementType": "$aScalar"
755
- },
756
- {
757
- "name": "norm_scale",
758
- "arg": "normScaleT",
759
- "semantic": "norm_scale",
760
- "buffer": { "type": "read-only-storage" },
761
- "elementType": "$aScalar",
762
- "length": "$K_LEN"
763
- },
764
- {
765
- "name": "gate_b",
766
- "arg": "gateBT",
767
- "semantic": "gate_B",
768
- "buffer": { "type": "read-only-storage" },
769
- "elementType": "u32"
770
- },
771
- {
772
- "name": "gate_scales",
773
- "arg": "gateScalesT",
774
- "semantic": "gate_scales",
775
- "buffer": { "type": "read-only-storage" },
776
- "elementType": "$aScalar"
777
- },
778
- {
779
- "name": "up_b",
780
- "arg": "upBT",
781
- "semantic": "up_B",
782
- "buffer": { "type": "read-only-storage" },
783
- "elementType": "u32"
784
- },
785
- {
786
- "name": "up_scales",
787
- "arg": "upScalesT",
788
- "semantic": "up_scales",
789
- "buffer": { "type": "read-only-storage" },
790
- "elementType": "$aScalar"
791
- },
792
- {
793
- "name": "up_bias",
794
- "arg": "upBiasT",
795
- "semantic": "up_bias",
796
- "buffer": { "type": "read-only-storage" },
797
- "elementType": "$aScalar",
798
- "length": "$N_LEN"
799
- },
800
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
801
- ],
802
- "fused_skip_nogb_ub": [
803
- {
804
- "name": "a",
805
- "arg": "aT",
806
- "semantic": "A",
807
- "buffer": { "type": "read-only-storage" },
808
- "elementType": "$aScalar"
809
- },
810
- {
811
- "name": "skip",
812
- "arg": "skipT",
813
- "semantic": "skip",
814
- "buffer": { "type": "read-only-storage" },
815
- "elementType": "$aScalar"
816
- },
817
- {
818
- "name": "norm_scale",
819
- "arg": "normScaleT",
820
- "semantic": "norm_scale",
821
- "buffer": { "type": "read-only-storage" },
822
- "elementType": "$aScalar",
823
- "length": "$K_LEN"
824
- },
825
- {
826
- "name": "gate_b",
827
- "arg": "gateBT",
828
- "semantic": "gate_B",
829
- "buffer": { "type": "read-only-storage" },
830
- "elementType": "u32"
831
- },
832
- {
833
- "name": "gate_scales",
834
- "arg": "gateScalesT",
835
- "semantic": "gate_scales",
836
- "buffer": { "type": "read-only-storage" },
837
- "elementType": "$aScalar"
838
- },
839
- {
840
- "name": "up_b",
841
- "arg": "upBT",
842
- "semantic": "up_B",
843
- "buffer": { "type": "read-only-storage" },
844
- "elementType": "u32"
845
- },
846
- {
847
- "name": "up_scales",
848
- "arg": "upScalesT",
849
- "semantic": "up_scales",
850
- "buffer": { "type": "read-only-storage" },
851
- "elementType": "$aScalar"
852
- },
853
- {
854
- "name": "up_bias",
855
- "arg": "upBiasT",
856
- "semantic": "up_bias",
857
- "buffer": { "type": "read-only-storage" },
858
- "elementType": "$aScalar",
859
- "length": "$N_LEN"
860
- },
861
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
862
- ],
863
- "fused_skipsum_nogb_ub": [
864
- {
865
- "name": "a",
866
- "arg": "aT",
867
- "semantic": "A",
868
- "buffer": { "type": "read-only-storage" },
869
- "elementType": "$aScalar"
870
- },
871
- {
872
- "name": "skip",
873
- "arg": "skipT",
874
- "semantic": "skip",
875
- "buffer": { "type": "read-only-storage" },
876
- "elementType": "$aScalar"
877
- },
878
- {
879
- "name": "norm_scale",
880
- "arg": "normScaleT",
881
- "semantic": "norm_scale",
882
- "buffer": { "type": "read-only-storage" },
883
- "elementType": "$aScalar",
884
- "length": "$K_LEN"
885
- },
886
- {
887
- "name": "gate_b",
888
- "arg": "gateBT",
889
- "semantic": "gate_B",
890
- "buffer": { "type": "read-only-storage" },
891
- "elementType": "u32"
892
- },
893
- {
894
- "name": "gate_scales",
895
- "arg": "gateScalesT",
896
- "semantic": "gate_scales",
897
- "buffer": { "type": "read-only-storage" },
898
- "elementType": "$aScalar"
899
- },
900
- {
901
- "name": "up_b",
902
- "arg": "upBT",
903
- "semantic": "up_B",
904
- "buffer": { "type": "read-only-storage" },
905
- "elementType": "u32"
906
- },
907
- {
908
- "name": "up_scales",
909
- "arg": "upScalesT",
910
- "semantic": "up_scales",
911
- "buffer": { "type": "read-only-storage" },
912
- "elementType": "$aScalar"
913
- },
914
- {
915
- "name": "up_bias",
916
- "arg": "upBiasT",
917
- "semantic": "up_bias",
918
- "buffer": { "type": "read-only-storage" },
919
- "elementType": "$aScalar",
920
- "length": "$N_LEN"
921
- },
922
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
923
- {
924
- "name": "residual",
925
- "arg": "residualT",
926
- "semantic": "input_skip_bias_sum",
927
- "buffer": { "type": "storage" },
928
- "elementType": "$aScalar"
929
- }
930
- ],
931
- "gateUpA_gb_noub": [
932
- {
933
- "name": "a",
934
- "arg": "aT",
935
- "semantic": "A",
936
- "buffer": { "type": "read-only-storage" },
937
- "elementType": "$aScalar"
938
- },
939
- {
940
- "name": "gate_b",
941
- "arg": "gateBT",
942
- "semantic": "gate_B",
943
- "buffer": { "type": "read-only-storage" },
944
- "elementType": "u32"
945
- },
946
- {
947
- "name": "gate_scales",
948
- "arg": "gateScalesT",
949
- "semantic": "gate_scales",
950
- "buffer": { "type": "read-only-storage" },
951
- "elementType": "$aScalar"
952
- },
953
- {
954
- "name": "gate_bias",
955
- "arg": "gateBiasT",
956
- "semantic": "gate_bias",
957
- "buffer": { "type": "read-only-storage" },
958
- "elementType": "$aScalar",
959
- "length": "$N_LEN"
960
- },
961
- {
962
- "name": "up_b",
963
- "arg": "upBT",
964
- "semantic": "up_B",
965
- "buffer": { "type": "read-only-storage" },
966
- "elementType": "u32"
967
- },
968
- {
969
- "name": "up_scales",
970
- "arg": "upScalesT",
971
- "semantic": "up_scales",
972
- "buffer": { "type": "read-only-storage" },
973
- "elementType": "$aScalar"
974
- },
975
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
976
- ],
977
- "gateUpNormed_gb_noub": [
978
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
979
- {
980
- "name": "gate_b",
981
- "arg": "gateBT",
982
- "semantic": "gate_B",
983
- "buffer": { "type": "read-only-storage" },
984
- "elementType": "u32"
985
- },
986
- {
987
- "name": "gate_scales",
988
- "arg": "gateScalesT",
989
- "semantic": "gate_scales",
990
- "buffer": { "type": "read-only-storage" },
991
- "elementType": "$aScalar"
992
- },
993
- {
994
- "name": "gate_bias",
995
- "arg": "gateBiasT",
996
- "semantic": "gate_bias",
997
- "buffer": { "type": "read-only-storage" },
998
- "elementType": "$aScalar",
999
- "length": "$N_LEN"
1000
- },
1001
- {
1002
- "name": "up_b",
1003
- "arg": "upBT",
1004
- "semantic": "up_B",
1005
- "buffer": { "type": "read-only-storage" },
1006
- "elementType": "u32"
1007
- },
1008
- {
1009
- "name": "up_scales",
1010
- "arg": "upScalesT",
1011
- "semantic": "up_scales",
1012
- "buffer": { "type": "read-only-storage" },
1013
- "elementType": "$aScalar"
1014
- },
1015
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1016
- ],
1017
- "fused_norm_gb_noub": [
1018
- {
1019
- "name": "a",
1020
- "arg": "aT",
1021
- "semantic": "A",
1022
- "buffer": { "type": "read-only-storage" },
1023
- "elementType": "$aScalar"
1024
- },
1025
- {
1026
- "name": "norm_scale",
1027
- "arg": "normScaleT",
1028
- "semantic": "norm_scale",
1029
- "buffer": { "type": "read-only-storage" },
1030
- "elementType": "$aScalar",
1031
- "length": "$K_LEN"
1032
- },
1033
- {
1034
- "name": "gate_b",
1035
- "arg": "gateBT",
1036
- "semantic": "gate_B",
1037
- "buffer": { "type": "read-only-storage" },
1038
- "elementType": "u32"
1039
- },
1040
- {
1041
- "name": "gate_scales",
1042
- "arg": "gateScalesT",
1043
- "semantic": "gate_scales",
1044
- "buffer": { "type": "read-only-storage" },
1045
- "elementType": "$aScalar"
1046
- },
1047
- {
1048
- "name": "gate_bias",
1049
- "arg": "gateBiasT",
1050
- "semantic": "gate_bias",
1051
- "buffer": { "type": "read-only-storage" },
1052
- "elementType": "$aScalar",
1053
- "length": "$N_LEN"
1054
- },
1055
- {
1056
- "name": "up_b",
1057
- "arg": "upBT",
1058
- "semantic": "up_B",
1059
- "buffer": { "type": "read-only-storage" },
1060
- "elementType": "u32"
1061
- },
1062
- {
1063
- "name": "up_scales",
1064
- "arg": "upScalesT",
1065
- "semantic": "up_scales",
1066
- "buffer": { "type": "read-only-storage" },
1067
- "elementType": "$aScalar"
1068
- },
1069
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1070
- ],
1071
- "fused_skip_gb_noub": [
1072
- {
1073
- "name": "a",
1074
- "arg": "aT",
1075
- "semantic": "A",
1076
- "buffer": { "type": "read-only-storage" },
1077
- "elementType": "$aScalar"
1078
- },
1079
- {
1080
- "name": "skip",
1081
- "arg": "skipT",
1082
- "semantic": "skip",
1083
- "buffer": { "type": "read-only-storage" },
1084
- "elementType": "$aScalar"
1085
- },
1086
- {
1087
- "name": "norm_scale",
1088
- "arg": "normScaleT",
1089
- "semantic": "norm_scale",
1090
- "buffer": { "type": "read-only-storage" },
1091
- "elementType": "$aScalar",
1092
- "length": "$K_LEN"
1093
- },
1094
- {
1095
- "name": "gate_b",
1096
- "arg": "gateBT",
1097
- "semantic": "gate_B",
1098
- "buffer": { "type": "read-only-storage" },
1099
- "elementType": "u32"
1100
- },
1101
- {
1102
- "name": "gate_scales",
1103
- "arg": "gateScalesT",
1104
- "semantic": "gate_scales",
1105
- "buffer": { "type": "read-only-storage" },
1106
- "elementType": "$aScalar"
1107
- },
1108
- {
1109
- "name": "gate_bias",
1110
- "arg": "gateBiasT",
1111
- "semantic": "gate_bias",
1112
- "buffer": { "type": "read-only-storage" },
1113
- "elementType": "$aScalar",
1114
- "length": "$N_LEN"
1115
- },
1116
- {
1117
- "name": "up_b",
1118
- "arg": "upBT",
1119
- "semantic": "up_B",
1120
- "buffer": { "type": "read-only-storage" },
1121
- "elementType": "u32"
1122
- },
1123
- {
1124
- "name": "up_scales",
1125
- "arg": "upScalesT",
1126
- "semantic": "up_scales",
1127
- "buffer": { "type": "read-only-storage" },
1128
- "elementType": "$aScalar"
1129
- },
1130
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1131
- ],
1132
- "fused_skipsum_gb_noub": [
1133
- {
1134
- "name": "a",
1135
- "arg": "aT",
1136
- "semantic": "A",
1137
- "buffer": { "type": "read-only-storage" },
1138
- "elementType": "$aScalar"
1139
- },
1140
- {
1141
- "name": "skip",
1142
- "arg": "skipT",
1143
- "semantic": "skip",
1144
- "buffer": { "type": "read-only-storage" },
1145
- "elementType": "$aScalar"
1146
- },
1147
- {
1148
- "name": "norm_scale",
1149
- "arg": "normScaleT",
1150
- "semantic": "norm_scale",
1151
- "buffer": { "type": "read-only-storage" },
1152
- "elementType": "$aScalar",
1153
- "length": "$K_LEN"
1154
- },
1155
- {
1156
- "name": "gate_b",
1157
- "arg": "gateBT",
1158
- "semantic": "gate_B",
1159
- "buffer": { "type": "read-only-storage" },
1160
- "elementType": "u32"
1161
- },
1162
- {
1163
- "name": "gate_scales",
1164
- "arg": "gateScalesT",
1165
- "semantic": "gate_scales",
1166
- "buffer": { "type": "read-only-storage" },
1167
- "elementType": "$aScalar"
1168
- },
1169
- {
1170
- "name": "gate_bias",
1171
- "arg": "gateBiasT",
1172
- "semantic": "gate_bias",
1173
- "buffer": { "type": "read-only-storage" },
1174
- "elementType": "$aScalar",
1175
- "length": "$N_LEN"
1176
- },
1177
- {
1178
- "name": "up_b",
1179
- "arg": "upBT",
1180
- "semantic": "up_B",
1181
- "buffer": { "type": "read-only-storage" },
1182
- "elementType": "u32"
1183
- },
1184
- {
1185
- "name": "up_scales",
1186
- "arg": "upScalesT",
1187
- "semantic": "up_scales",
1188
- "buffer": { "type": "read-only-storage" },
1189
- "elementType": "$aScalar"
1190
- },
1191
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
1192
- {
1193
- "name": "residual",
1194
- "arg": "residualT",
1195
- "semantic": "input_skip_bias_sum",
1196
- "buffer": { "type": "storage" },
1197
- "elementType": "$aScalar"
1198
- }
1199
- ],
1200
- "gateUpA_gb_ub": [
1201
- {
1202
- "name": "a",
1203
- "arg": "aT",
1204
- "semantic": "A",
1205
- "buffer": { "type": "read-only-storage" },
1206
- "elementType": "$aScalar"
1207
- },
1208
- {
1209
- "name": "gate_b",
1210
- "arg": "gateBT",
1211
- "semantic": "gate_B",
1212
- "buffer": { "type": "read-only-storage" },
1213
- "elementType": "u32"
1214
- },
1215
- {
1216
- "name": "gate_scales",
1217
- "arg": "gateScalesT",
1218
- "semantic": "gate_scales",
1219
- "buffer": { "type": "read-only-storage" },
1220
- "elementType": "$aScalar"
1221
- },
1222
- {
1223
- "name": "gate_bias",
1224
- "arg": "gateBiasT",
1225
- "semantic": "gate_bias",
1226
- "buffer": { "type": "read-only-storage" },
1227
- "elementType": "$aScalar",
1228
- "length": "$N_LEN"
1229
- },
1230
- {
1231
- "name": "up_b",
1232
- "arg": "upBT",
1233
- "semantic": "up_B",
1234
- "buffer": { "type": "read-only-storage" },
1235
- "elementType": "u32"
1236
- },
1237
- {
1238
- "name": "up_scales",
1239
- "arg": "upScalesT",
1240
- "semantic": "up_scales",
1241
- "buffer": { "type": "read-only-storage" },
1242
- "elementType": "$aScalar"
1243
- },
1244
- {
1245
- "name": "up_bias",
1246
- "arg": "upBiasT",
1247
- "semantic": "up_bias",
1248
- "buffer": { "type": "read-only-storage" },
1249
- "elementType": "$aScalar",
1250
- "length": "$N_LEN"
1251
- },
1252
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1253
- ],
1254
- "gateUpNormed_gb_ub": [
1255
- { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
1256
- {
1257
- "name": "gate_b",
1258
- "arg": "gateBT",
1259
- "semantic": "gate_B",
1260
- "buffer": { "type": "read-only-storage" },
1261
- "elementType": "u32"
1262
- },
1263
- {
1264
- "name": "gate_scales",
1265
- "arg": "gateScalesT",
1266
- "semantic": "gate_scales",
1267
- "buffer": { "type": "read-only-storage" },
1268
- "elementType": "$aScalar"
1269
- },
1270
- {
1271
- "name": "gate_bias",
1272
- "arg": "gateBiasT",
1273
- "semantic": "gate_bias",
1274
- "buffer": { "type": "read-only-storage" },
1275
- "elementType": "$aScalar",
1276
- "length": "$N_LEN"
1277
- },
1278
- {
1279
- "name": "up_b",
1280
- "arg": "upBT",
1281
- "semantic": "up_B",
1282
- "buffer": { "type": "read-only-storage" },
1283
- "elementType": "u32"
1284
- },
1285
- {
1286
- "name": "up_scales",
1287
- "arg": "upScalesT",
1288
- "semantic": "up_scales",
1289
- "buffer": { "type": "read-only-storage" },
1290
- "elementType": "$aScalar"
1291
- },
1292
- {
1293
- "name": "up_bias",
1294
- "arg": "upBiasT",
1295
- "semantic": "up_bias",
1296
- "buffer": { "type": "read-only-storage" },
1297
- "elementType": "$aScalar",
1298
- "length": "$N_LEN"
1299
- },
1300
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1301
- ],
1302
- "fused_norm_gb_ub": [
1303
- {
1304
- "name": "a",
1305
- "arg": "aT",
1306
- "semantic": "A",
1307
- "buffer": { "type": "read-only-storage" },
1308
- "elementType": "$aScalar"
1309
- },
1310
- {
1311
- "name": "norm_scale",
1312
- "arg": "normScaleT",
1313
- "semantic": "norm_scale",
1314
- "buffer": { "type": "read-only-storage" },
1315
- "elementType": "$aScalar",
1316
- "length": "$K_LEN"
1317
- },
1318
- {
1319
- "name": "gate_b",
1320
- "arg": "gateBT",
1321
- "semantic": "gate_B",
1322
- "buffer": { "type": "read-only-storage" },
1323
- "elementType": "u32"
1324
- },
1325
- {
1326
- "name": "gate_scales",
1327
- "arg": "gateScalesT",
1328
- "semantic": "gate_scales",
1329
- "buffer": { "type": "read-only-storage" },
1330
- "elementType": "$aScalar"
1331
- },
1332
- {
1333
- "name": "gate_bias",
1334
- "arg": "gateBiasT",
1335
- "semantic": "gate_bias",
1336
- "buffer": { "type": "read-only-storage" },
1337
- "elementType": "$aScalar",
1338
- "length": "$N_LEN"
1339
- },
1340
- {
1341
- "name": "up_b",
1342
- "arg": "upBT",
1343
- "semantic": "up_B",
1344
- "buffer": { "type": "read-only-storage" },
1345
- "elementType": "u32"
1346
- },
1347
- {
1348
- "name": "up_scales",
1349
- "arg": "upScalesT",
1350
- "semantic": "up_scales",
1351
- "buffer": { "type": "read-only-storage" },
1352
- "elementType": "$aScalar"
1353
- },
1354
- {
1355
- "name": "up_bias",
1356
- "arg": "upBiasT",
1357
- "semantic": "up_bias",
1358
- "buffer": { "type": "read-only-storage" },
1359
- "elementType": "$aScalar",
1360
- "length": "$N_LEN"
1361
- },
1362
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1363
- ],
1364
- "fused_skip_gb_ub": [
1365
- {
1366
- "name": "a",
1367
- "arg": "aT",
1368
- "semantic": "A",
1369
- "buffer": { "type": "read-only-storage" },
1370
- "elementType": "$aScalar"
1371
- },
1372
- {
1373
- "name": "skip",
1374
- "arg": "skipT",
1375
- "semantic": "skip",
1376
- "buffer": { "type": "read-only-storage" },
1377
- "elementType": "$aScalar"
1378
- },
1379
- {
1380
- "name": "norm_scale",
1381
- "arg": "normScaleT",
1382
- "semantic": "norm_scale",
1383
- "buffer": { "type": "read-only-storage" },
1384
- "elementType": "$aScalar",
1385
- "length": "$K_LEN"
1386
- },
1387
- {
1388
- "name": "gate_b",
1389
- "arg": "gateBT",
1390
- "semantic": "gate_B",
1391
- "buffer": { "type": "read-only-storage" },
1392
- "elementType": "u32"
1393
- },
1394
- {
1395
- "name": "gate_scales",
1396
- "arg": "gateScalesT",
1397
- "semantic": "gate_scales",
1398
- "buffer": { "type": "read-only-storage" },
1399
- "elementType": "$aScalar"
1400
- },
1401
- {
1402
- "name": "gate_bias",
1403
- "arg": "gateBiasT",
1404
- "semantic": "gate_bias",
1405
- "buffer": { "type": "read-only-storage" },
1406
- "elementType": "$aScalar",
1407
- "length": "$N_LEN"
1408
- },
1409
- {
1410
- "name": "up_b",
1411
- "arg": "upBT",
1412
- "semantic": "up_B",
1413
- "buffer": { "type": "read-only-storage" },
1414
- "elementType": "u32"
1415
- },
1416
- {
1417
- "name": "up_scales",
1418
- "arg": "upScalesT",
1419
- "semantic": "up_scales",
1420
- "buffer": { "type": "read-only-storage" },
1421
- "elementType": "$aScalar"
1422
- },
1423
- {
1424
- "name": "up_bias",
1425
- "arg": "upBiasT",
1426
- "semantic": "up_bias",
1427
- "buffer": { "type": "read-only-storage" },
1428
- "elementType": "$aScalar",
1429
- "length": "$N_LEN"
1430
- },
1431
- { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
1432
- ]
1433
  },
1434
  "variants": [
1435
  {
1436
  "id": "plain_nogb_noub",
1437
  "priority": 10,
1438
- "when": ["mlpShapeOk", "not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_noub", "gateUpDispatchFits"],
1439
- "constants": { "inlineNorm": "0", "fromNormed": "0" },
1440
  "passes": [
1441
  {
1442
  "id": "main",
1443
  "name": "MatMulNBitsMlp.GateUp",
1444
  "shader": "mlp-gate-up.wgsl.jinja",
1445
- "bindings": "gateUpA_nogb_noub",
1446
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
1447
  }
1448
  ]
1449
  },
1450
  {
1451
  "id": "plain_nogb_ub",
1452
  "priority": 10,
1453
- "when": ["mlpShapeOk", "not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_ub", "gateUpDispatchFits"],
1454
- "constants": { "inlineNorm": "0", "fromNormed": "0" },
1455
  "passes": [
1456
  {
1457
  "id": "main",
1458
  "name": "MatMulNBitsMlp.GateUp",
1459
  "shader": "mlp-gate-up.wgsl.jinja",
1460
- "bindings": "gateUpA_nogb_ub",
1461
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
1462
  }
1463
  ]
1464
  },
1465
  {
1466
  "id": "plain_gb_noub",
1467
  "priority": 10,
1468
- "when": ["mlpShapeOk", "not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_noub", "gateUpDispatchFits"],
1469
- "constants": { "inlineNorm": "0", "fromNormed": "0" },
1470
  "passes": [
1471
  {
1472
  "id": "main",
1473
  "name": "MatMulNBitsMlp.GateUp",
1474
  "shader": "mlp-gate-up.wgsl.jinja",
1475
- "bindings": "gateUpA_gb_noub",
1476
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
1477
  }
1478
  ]
1479
  },
1480
  {
1481
  "id": "plain_gb_ub",
1482
  "priority": 10,
1483
- "when": ["mlpShapeOk", "not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_ub", "gateUpDispatchFits"],
1484
- "constants": { "inlineNorm": "0", "fromNormed": "0" },
1485
  "passes": [
1486
  {
1487
  "id": "main",
1488
  "name": "MatMulNBitsMlp.GateUp",
1489
  "shader": "mlp-gate-up.wgsl.jinja",
1490
- "bindings": "gateUpA_gb_ub",
1491
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1492
  }
1493
  ]
1494
  },
1495
  {
1496
  "id": "staged_norm_nogb_noub",
1497
  "priority": 10,
1498
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_noub", "gateUpDispatchFits", "normDispatchFits", "not present.skipT", "not present.residualT"],
1499
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1500
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1501
  "passes": [
1502
  {
1503
  "id": "norm",
1504
  "name": "MatMulNBitsMlp.RmsNorm",
1505
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1506
- "bindings": "normOnly",
1507
- "dispatch": { "workgroups": "aRows" }
1508
  },
1509
  {
1510
  "id": "main",
1511
  "name": "MatMulNBitsMlp.GateUp",
1512
  "shader": "mlp-gate-up.wgsl.jinja",
1513
- "bindings": "gateUpNormed_nogb_noub",
1514
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1515
  }
1516
  ]
1517
  },
1518
  {
1519
  "id": "staged_skip_nogb_noub",
1520
  "priority": 10,
1521
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_noub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "not present.residualT"],
1522
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1523
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1524
  "passes": [
1525
  {
1526
  "id": "norm",
1527
  "name": "MatMulNBitsMlp.RmsNorm",
1528
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1529
- "bindings": "normSkip",
1530
- "dispatch": { "workgroups": "aRows" }
1531
  },
1532
  {
1533
  "id": "main",
1534
  "name": "MatMulNBitsMlp.GateUp",
1535
  "shader": "mlp-gate-up.wgsl.jinja",
1536
- "bindings": "gateUpNormed_nogb_noub",
1537
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1538
  }
1539
  ]
1540
  },
1541
  {
1542
  "id": "staged_skipsum_nogb_noub",
1543
  "priority": 10,
1544
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_noub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "present.residualT"],
1545
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1546
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1547
  "passes": [
1548
  {
1549
  "id": "norm",
1550
  "name": "MatMulNBitsMlp.RmsNorm",
1551
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1552
- "bindings": "normFull",
1553
- "dispatch": { "workgroups": "aRows" }
1554
  },
1555
  {
1556
  "id": "main",
1557
  "name": "MatMulNBitsMlp.GateUp",
1558
  "shader": "mlp-gate-up.wgsl.jinja",
1559
- "bindings": "gateUpNormed_nogb_noub",
1560
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1561
  }
1562
  ]
1563
  },
1564
  {
1565
  "id": "staged_norm_nogb_ub",
1566
  "priority": 10,
1567
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_ub", "gateUpDispatchFits", "normDispatchFits", "not present.skipT", "not present.residualT"],
1568
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1569
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1570
  "passes": [
1571
  {
1572
  "id": "norm",
1573
  "name": "MatMulNBitsMlp.RmsNorm",
1574
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1575
- "bindings": "normOnly",
1576
- "dispatch": { "workgroups": "aRows" }
1577
  },
1578
  {
1579
  "id": "main",
1580
  "name": "MatMulNBitsMlp.GateUp",
1581
  "shader": "mlp-gate-up.wgsl.jinja",
1582
- "bindings": "gateUpNormed_nogb_ub",
1583
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1584
  }
1585
  ]
1586
  },
1587
  {
1588
  "id": "staged_skip_nogb_ub",
1589
  "priority": 10,
1590
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_ub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "not present.residualT"],
1591
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1592
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1593
  "passes": [
1594
  {
1595
  "id": "norm",
1596
  "name": "MatMulNBitsMlp.RmsNorm",
1597
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1598
- "bindings": "normSkip",
1599
- "dispatch": { "workgroups": "aRows" }
1600
  },
1601
  {
1602
  "id": "main",
1603
  "name": "MatMulNBitsMlp.GateUp",
1604
  "shader": "mlp-gate-up.wgsl.jinja",
1605
- "bindings": "gateUpNormed_nogb_ub",
1606
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1607
  }
1608
  ]
1609
  },
1610
  {
1611
  "id": "staged_skipsum_nogb_ub",
1612
  "priority": 10,
1613
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_ub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "present.residualT"],
1614
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1615
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1616
  "passes": [
1617
  {
1618
  "id": "norm",
1619
  "name": "MatMulNBitsMlp.RmsNorm",
1620
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1621
- "bindings": "normFull",
1622
- "dispatch": { "workgroups": "aRows" }
1623
  },
1624
  {
1625
  "id": "main",
1626
  "name": "MatMulNBitsMlp.GateUp",
1627
  "shader": "mlp-gate-up.wgsl.jinja",
1628
- "bindings": "gateUpNormed_nogb_ub",
1629
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1630
  }
1631
  ]
1632
  },
1633
  {
1634
  "id": "staged_norm_gb_noub",
1635
  "priority": 10,
1636
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_noub", "gateUpDispatchFits", "normDispatchFits", "not present.skipT", "not present.residualT"],
1637
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1638
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1639
  "passes": [
1640
  {
1641
  "id": "norm",
1642
  "name": "MatMulNBitsMlp.RmsNorm",
1643
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1644
- "bindings": "normOnly",
1645
- "dispatch": { "workgroups": "aRows" }
1646
  },
1647
  {
1648
  "id": "main",
1649
  "name": "MatMulNBitsMlp.GateUp",
1650
  "shader": "mlp-gate-up.wgsl.jinja",
1651
- "bindings": "gateUpNormed_gb_noub",
1652
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1653
  }
1654
  ]
1655
  },
1656
  {
1657
  "id": "staged_skip_gb_noub",
1658
  "priority": 10,
1659
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_noub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "not present.residualT"],
1660
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1661
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1662
  "passes": [
1663
  {
1664
  "id": "norm",
1665
  "name": "MatMulNBitsMlp.RmsNorm",
1666
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1667
- "bindings": "normSkip",
1668
- "dispatch": { "workgroups": "aRows" }
1669
  },
1670
  {
1671
  "id": "main",
1672
  "name": "MatMulNBitsMlp.GateUp",
1673
  "shader": "mlp-gate-up.wgsl.jinja",
1674
- "bindings": "gateUpNormed_gb_noub",
1675
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1676
  }
1677
  ]
1678
  },
1679
  {
1680
  "id": "staged_skipsum_gb_noub",
1681
  "priority": 10,
1682
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_noub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "present.residualT"],
1683
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1684
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1685
  "passes": [
1686
  {
1687
  "id": "norm",
1688
  "name": "MatMulNBitsMlp.RmsNorm",
1689
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1690
- "bindings": "normFull",
1691
- "dispatch": { "workgroups": "aRows" }
1692
  },
1693
  {
1694
  "id": "main",
1695
  "name": "MatMulNBitsMlp.GateUp",
1696
  "shader": "mlp-gate-up.wgsl.jinja",
1697
- "bindings": "gateUpNormed_gb_noub",
1698
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1699
  }
1700
  ]
1701
  },
1702
  {
1703
  "id": "staged_norm_gb_ub",
1704
  "priority": 10,
1705
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_ub", "gateUpDispatchFits", "normDispatchFits", "not present.skipT", "not present.residualT"],
1706
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1707
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1708
  "passes": [
1709
  {
1710
  "id": "norm",
1711
  "name": "MatMulNBitsMlp.RmsNorm",
1712
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1713
- "bindings": "normOnly",
1714
- "dispatch": { "workgroups": "aRows" }
1715
  },
1716
  {
1717
  "id": "main",
1718
  "name": "MatMulNBitsMlp.GateUp",
1719
  "shader": "mlp-gate-up.wgsl.jinja",
1720
- "bindings": "gateUpNormed_gb_ub",
1721
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1722
  }
1723
  ]
1724
  },
1725
  {
1726
  "id": "staged_skip_gb_ub",
1727
  "priority": 10,
1728
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_ub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "not present.residualT"],
1729
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1730
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1731
  "passes": [
1732
  {
1733
  "id": "norm",
1734
  "name": "MatMulNBitsMlp.RmsNorm",
1735
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1736
- "bindings": "normSkip",
1737
- "dispatch": { "workgroups": "aRows" }
1738
  },
1739
  {
1740
  "id": "main",
1741
  "name": "MatMulNBitsMlp.GateUp",
1742
  "shader": "mlp-gate-up.wgsl.jinja",
1743
- "bindings": "gateUpNormed_gb_ub",
1744
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1745
  }
1746
  ]
1747
  },
1748
  {
1749
  "id": "staged_skipsum_gb_ub",
1750
  "priority": 10,
1751
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_ub", "gateUpDispatchFits", "normDispatchFits", "present.skipT", "present.residualT"],
1752
- "constants": { "inlineNorm": "0", "fromNormed": "1" },
1753
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
1754
  "passes": [
1755
  {
1756
  "id": "norm",
1757
  "name": "MatMulNBitsMlp.RmsNorm",
1758
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
1759
- "bindings": "normFull",
1760
- "dispatch": { "workgroups": "aRows" }
1761
  },
1762
  {
1763
  "id": "main",
1764
  "name": "MatMulNBitsMlp.GateUp",
1765
  "shader": "mlp-gate-up.wgsl.jinja",
1766
- "bindings": "gateUpNormed_gb_ub",
1767
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "rowGroups" }
 
1768
  }
1769
  ]
1770
  },
1771
  {
1772
  "id": "fused_norm_nogb_noub",
1773
  "priority": 30,
 
1774
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 7 } },
1775
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_noub", "gateUpDispatchFits", "aRows == 1", "not present.skipT", "not present.residualT"],
1776
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1777
  "passes": [
1778
  {
1779
  "id": "main",
1780
  "name": "MatMulNBitsMlp.FusedDecode",
1781
  "shader": "mlp-gate-up.wgsl.jinja",
1782
- "bindings": "fused_norm_nogb_noub",
1783
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1784
  }
1785
  ]
1786
  },
1787
  {
1788
  "id": "fused_skip_nogb_noub",
1789
  "priority": 30,
 
1790
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
1791
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_noub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "not present.residualT"],
1792
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1793
  "passes": [
1794
  {
1795
  "id": "main",
1796
  "name": "MatMulNBitsMlp.FusedDecode",
1797
  "shader": "mlp-gate-up.wgsl.jinja",
1798
- "bindings": "fused_skip_nogb_noub",
1799
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1800
  }
1801
  ]
1802
  },
1803
  {
1804
  "id": "fused_skipsum_nogb_noub",
1805
  "priority": 30,
 
1806
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
1807
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_noub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "present.residualT"],
1808
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1809
  "passes": [
1810
  {
1811
  "id": "main",
1812
  "name": "MatMulNBitsMlp.FusedDecode",
1813
  "shader": "mlp-gate-up.wgsl.jinja",
1814
- "bindings": "fused_skipsum_nogb_noub",
1815
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1816
  }
1817
  ]
1818
  },
1819
  {
1820
  "id": "fused_norm_nogb_ub",
1821
  "priority": 30,
 
1822
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
1823
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_ub", "gateUpDispatchFits", "aRows == 1", "not present.skipT", "not present.residualT"],
1824
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1825
  "passes": [
1826
  {
1827
  "id": "main",
1828
  "name": "MatMulNBitsMlp.FusedDecode",
1829
  "shader": "mlp-gate-up.wgsl.jinja",
1830
- "bindings": "fused_norm_nogb_ub",
1831
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1832
  }
1833
  ]
1834
  },
1835
  {
1836
  "id": "fused_skip_nogb_ub",
1837
  "priority": 30,
 
1838
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
1839
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_ub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "not present.residualT"],
1840
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1841
  "passes": [
1842
  {
1843
  "id": "main",
1844
  "name": "MatMulNBitsMlp.FusedDecode",
1845
  "shader": "mlp-gate-up.wgsl.jinja",
1846
- "bindings": "fused_skip_nogb_ub",
1847
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1848
  }
1849
  ]
1850
  },
1851
  {
1852
  "id": "fused_skipsum_nogb_ub",
1853
  "priority": 30,
 
1854
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
1855
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_nogb_ub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "present.residualT"],
1856
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1857
  "passes": [
1858
  {
1859
  "id": "main",
1860
  "name": "MatMulNBitsMlp.FusedDecode",
1861
  "shader": "mlp-gate-up.wgsl.jinja",
1862
- "bindings": "fused_skipsum_nogb_ub",
1863
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1864
  }
1865
  ]
1866
  },
1867
  {
1868
  "id": "fused_norm_gb_noub",
1869
  "priority": 30,
 
1870
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
1871
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_noub", "gateUpDispatchFits", "aRows == 1", "not present.skipT", "not present.residualT"],
1872
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1873
  "passes": [
1874
  {
1875
  "id": "main",
1876
  "name": "MatMulNBitsMlp.FusedDecode",
1877
  "shader": "mlp-gate-up.wgsl.jinja",
1878
- "bindings": "fused_norm_gb_noub",
1879
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1880
  }
1881
  ]
1882
  },
1883
  {
1884
  "id": "fused_skip_gb_noub",
1885
  "priority": 30,
 
1886
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
1887
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_noub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "not present.residualT"],
1888
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1889
  "passes": [
1890
  {
1891
  "id": "main",
1892
  "name": "MatMulNBitsMlp.FusedDecode",
1893
  "shader": "mlp-gate-up.wgsl.jinja",
1894
- "bindings": "fused_skip_gb_noub",
1895
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1896
  }
1897
  ]
1898
  },
1899
  {
1900
  "id": "fused_skipsum_gb_noub",
1901
  "priority": 30,
 
1902
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
1903
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_noub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "present.residualT"],
1904
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1905
  "passes": [
1906
  {
1907
  "id": "main",
1908
  "name": "MatMulNBitsMlp.FusedDecode",
1909
  "shader": "mlp-gate-up.wgsl.jinja",
1910
- "bindings": "fused_skipsum_gb_noub",
1911
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1912
  }
1913
  ]
1914
  },
1915
  {
1916
  "id": "fused_norm_gb_ub",
1917
  "priority": 30,
 
1918
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
1919
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_ub", "gateUpDispatchFits", "aRows == 1", "not present.skipT", "not present.residualT"],
1920
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1921
  "passes": [
1922
  {
1923
  "id": "main",
1924
  "name": "MatMulNBitsMlp.FusedDecode",
1925
  "shader": "mlp-gate-up.wgsl.jinja",
1926
- "bindings": "fused_norm_gb_ub",
1927
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1928
  }
1929
  ]
1930
  },
1931
  {
1932
  "id": "fused_skip_gb_ub",
1933
  "priority": 30,
 
1934
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
1935
- "when": ["mlpShapeOk", "normContractOk", "biasPresence_gb_ub", "gateUpDispatchFits", "aRows == 1", "present.skipT", "not present.residualT"],
1936
- "constants": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
1937
  "passes": [
1938
  {
1939
  "id": "main",
1940
  "name": "MatMulNBitsMlp.FusedDecode",
1941
  "shader": "mlp-gate-up.wgsl.jinja",
1942
- "bindings": "fused_skip_gb_ub",
1943
- "dispatch": { "x": "ceilDiv(attrs.N, tunables.TILE_N)", "y": "aRows" }
 
1944
  }
1945
  ]
1946
  }
 
2
  "domain": "com.microsoft",
3
  "name": "MatMulNBitsMlp",
4
  "sinceVersion": 1,
5
+ "inputs": {
6
+ "aT": { "onnx": "A", "dtype": "T1" },
7
+ "skipT": { "onnx": "skip", "dtype": "T1", "optional": true },
8
+ "normScaleT": { "onnx": "norm_scale", "dtype": "T1", "rank": 1, "optional": true },
9
+ "gateBT": { "onnx": "gate_B", "dtype": "uint8", "rank": 3, "layout": "packed" },
10
+ "gateScalesT": { "onnx": "gate_scales", "dtype": "T1", "rank": 2 },
11
+ "gateBiasT": { "onnx": "gate_bias", "dtype": "T1", "rank": 1, "optional": true },
12
+ "upBT": { "onnx": "up_B", "dtype": "uint8", "rank": 3, "layout": "packed" },
13
+ "upScalesT": { "onnx": "up_scales", "dtype": "T1", "rank": 2 },
14
+ "upBiasT": { "onnx": "up_bias", "dtype": "T1", "rank": 1, "optional": true }
15
+ },
16
+ "outputs": {
17
+ "yT": { "onnx": "Y", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.N]" },
18
+ "residualT": {
19
+ "onnx": "input_skip_bias_sum",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
  "dtype": "T1",
21
  "rank": "ranks.aT",
22
  "optional": true,
23
+ "shape": "shapes.aT"
 
24
  }
25
+ },
26
+ "attributes": {
27
+ "accuracy_level": { "default": 0 },
28
+ "bits": { "default": 4 },
29
+ "epsilon": { "default": 0.00001 },
30
+ "K": {},
31
+ "N": {},
32
+ "activation": {},
33
+ "block_size": {}
 
34
  },
35
  "attributeConstraints": {
36
  "K": { "required": true },
 
41
  "block_size": { "required": true }
42
  },
43
  "typeConstraints": { "T1": ["float32", "float16"] },
44
+ "tunables": {
45
+ "TILE_N": { "default": 8 },
46
+ "LANES": { "default": 8 },
47
+ "NORM_WORKGROUP_SIZE": { "default": 128 },
48
+ "ROW_TILE": { "default": 8 },
49
+ "DECODE_WORKGROUP_SIZE": { "default": 64 }
 
 
 
 
 
 
50
  },
 
51
  "derive": {
52
  "aRows": "numel(shapes.aT) / max(1, attrs.K)",
53
+ "rowTilePlan": "1 if aRows <= 1 else min(aRows, tunables.ROW_TILE)",
54
+ "rowGroups": "ceilDiv(aRows, rowTilePlan)",
55
  "kBlocks": "dim(shapes.gateBT, 1)",
56
  "blobSize": "dim(shapes.gateBT, 2)",
57
  "codesPerByte": "8 / attrs.bits",
 
66
  "lanesPow2": "tunables.LANES == pow2ceil(tunables.LANES)",
67
  "mlpShapeOk": "bitsSupported and weightShapeOk and scaleShapeOk and ioShapeOk and biasShapeOk and dtypeOk and lanesPow2 and attrs.K > 0 and attrs.N > 0 and attrs.block_size > 0",
68
  "normContractOk": "present.normScaleT and ranks.normScaleT == 1 and dim(shapes.normScaleT, 0) == attrs.K and tensorDtypes.normScaleT == tensorDtypes.aT and (sameShape(shapes.skipT, shapes.aT) and tensorDtypes.skipT == tensorDtypes.aT if present.skipT else true) and (sameShape(shapes.residualT, shapes.aT) and tensorDtypes.residualT == tensorDtypes.aT and present.skipT if present.residualT else true)",
69
+ "decodeWalk": "aRows <= 1",
70
+ "decodeVec": "decodeWalk and blobSize % 16 == 0",
71
+ "decodeActVec4": "decodeVec and attrs.K % attrs.block_size == 0",
72
+ "decodeLaneSplit": "decodeVec and kBlocks * blobSize <= tunables.DECODE_WORKGROUP_SIZE * 16",
73
+ "decodeCols": "8 if decodeLaneSplit else 4",
74
+ "decodeLanes": "tunables.DECODE_WORKGROUP_SIZE * (2 if decodeLaneSplit else 1)",
75
+ "tileCols": "decodeCols if decodeWalk else tunables.TILE_N",
76
+ "decodeWorkgroupOk": "tunables.DECODE_WORKGROUP_SIZE >= 4 and pow2ceil(tunables.DECODE_WORKGROUP_SIZE) == tunables.DECODE_WORKGROUP_SIZE and decodeLanes <= device.limits.maxComputeInvocationsPerWorkgroup and decodeLanes <= device.limits.maxComputeWorkgroupSizeX",
77
+ "gateUpDispatchFits": "decodeWorkgroupOk and ceilDiv(attrs.N, tileCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and aRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX",
78
  "normDispatchFits": "tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX",
79
  "biasPresence_nogb_noub": "not present.gateBiasT and not present.upBiasT",
80
  "biasPresence_nogb_ub": "not present.gateBiasT and present.upBiasT",
81
  "biasPresence_gb_noub": "present.gateBiasT and not present.upBiasT",
82
+ "biasPresence_gb_ub": "present.gateBiasT and present.upBiasT",
 
 
83
  "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
84
  "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
 
85
  "K": "attrs.K",
86
  "N": "attrs.N",
87
  "blockSize": "attrs.block_size",
 
 
88
  "bits": "attrs.bits",
 
 
89
  "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
90
  "tileN": "tunables.TILE_N",
91
  "lanes": "tunables.LANES",
92
+ "rowTile": "rowTilePlan",
93
+ "rowCount": "aRows",
94
+ "decodeNCols": "decodeCols",
95
+ "decodeWorkgroupSize": "decodeLanes",
96
+ "laneGroups": "2 if decodeLaneSplit else 1",
97
+ "useSubgroups": "device.features.has(\"subgroups\")",
98
+ "weightElement": "\"vec4<u32>\" if decodeVec else \"u32\"",
99
+ "actVec4": "decodeActVec4",
100
+ "normedElement": "\"vec4<f32>\" if decodeActVec4 else \"f32\"",
101
  "hidden": "attrs.K",
102
  "workgroupSize": "tunables.NORM_WORKGROUP_SIZE",
103
  "epsilon": "epsilonValue",
 
108
  "K_LEN": "attrs.K",
109
  "N_LEN": "attrs.N"
110
  },
111
+ "when": ["mlpShapeOk", "gateUpDispatchFits"],
112
+ "bindings": {
113
+ "a": { "arg": "aT", "buffer": "read-only-storage", "elementType": "$aScalar" },
114
+ "gate_b": { "arg": "gateBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
115
+ "gate_scales": { "arg": "gateScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
116
+ "up_b": { "arg": "upBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
117
+ "up_scales": { "arg": "upScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
118
+ "y": { "arg": "yT", "buffer": "storage", "elementType": "$aScalar" },
119
+ "up_bias": { "arg": "upBiasT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$N_LEN" },
120
+ "gate_bias": { "arg": "gateBiasT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$N_LEN" },
121
+ "norm_scale": { "arg": "normScaleT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$K_LEN" },
122
+ "normed": { "scratch": "normedA", "buffer": "storage", "elementType": "f32" },
123
+ "params": { "buffer": "uniform", "struct": [{ "name": "rows", "type": "u32", "value": "aRows" }] },
124
+ "normed_2": {
125
+ "scratch": "normedA",
126
+ "name": "normed",
127
+ "buffer": "read-only-storage",
128
+ "elementType": "$normedElement"
129
+ },
130
+ "skip": { "arg": "skipT", "buffer": "read-only-storage", "elementType": "$aScalar" },
131
+ "residual": { "arg": "residualT", "buffer": "storage", "elementType": "$aScalar" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
132
  },
133
  "variants": [
134
  {
135
  "id": "plain_nogb_noub",
136
  "priority": 10,
137
+ "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_noub"],
138
+ "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
139
  "passes": [
140
  {
141
  "id": "main",
142
  "name": "MatMulNBitsMlp.GateUp",
143
  "shader": "mlp-gate-up.wgsl.jinja",
144
+ "bindings": ["a", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
145
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
146
  }
147
  ]
148
  },
149
  {
150
  "id": "plain_nogb_ub",
151
  "priority": 10,
152
+ "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_ub"],
153
+ "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
154
  "passes": [
155
  {
156
  "id": "main",
157
  "name": "MatMulNBitsMlp.GateUp",
158
  "shader": "mlp-gate-up.wgsl.jinja",
159
+ "bindings": ["a", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
160
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
161
  }
162
  ]
163
  },
164
  {
165
  "id": "plain_gb_noub",
166
  "priority": 10,
167
+ "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_noub"],
168
+ "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
169
  "passes": [
170
  {
171
  "id": "main",
172
  "name": "MatMulNBitsMlp.GateUp",
173
  "shader": "mlp-gate-up.wgsl.jinja",
174
+ "bindings": ["a", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
175
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
176
  }
177
  ]
178
  },
179
  {
180
  "id": "plain_gb_ub",
181
  "priority": 10,
182
+ "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_ub"],
183
+ "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
184
  "passes": [
185
  {
186
  "id": "main",
187
  "name": "MatMulNBitsMlp.GateUp",
188
  "shader": "mlp-gate-up.wgsl.jinja",
189
+ "bindings": ["a", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
190
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
191
+ "subgroupCollectivesWidth": "portable"
192
  }
193
  ]
194
  },
195
  {
196
  "id": "staged_norm_nogb_noub",
197
  "priority": 10,
198
+ "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "not present.skipT", "not present.residualT"],
199
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
200
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
201
  "passes": [
202
  {
203
  "id": "norm",
204
  "name": "MatMulNBitsMlp.RmsNorm",
205
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
206
+ "bindings": ["a", "norm_scale", "normed", "params"],
207
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
208
  },
209
  {
210
  "id": "main",
211
  "name": "MatMulNBitsMlp.GateUp",
212
  "shader": "mlp-gate-up.wgsl.jinja",
213
+ "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
214
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
215
+ "subgroupCollectivesWidth": "portable"
216
  }
217
  ]
218
  },
219
  {
220
  "id": "staged_skip_nogb_noub",
221
  "priority": 10,
222
+ "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "present.skipT", "not present.residualT"],
223
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
224
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
225
  "passes": [
226
  {
227
  "id": "norm",
228
  "name": "MatMulNBitsMlp.RmsNorm",
229
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
230
+ "bindings": ["a", "skip", "norm_scale", "normed", "params"],
231
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
232
  },
233
  {
234
  "id": "main",
235
  "name": "MatMulNBitsMlp.GateUp",
236
  "shader": "mlp-gate-up.wgsl.jinja",
237
+ "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
238
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
239
+ "subgroupCollectivesWidth": "portable"
240
  }
241
  ]
242
  },
243
  {
244
  "id": "staged_skipsum_nogb_noub",
245
  "priority": 10,
246
+ "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "present.skipT", "present.residualT"],
247
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
248
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
249
  "passes": [
250
  {
251
  "id": "norm",
252
  "name": "MatMulNBitsMlp.RmsNorm",
253
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
254
+ "bindings": ["a", "skip", "norm_scale", "normed", "residual", "params"],
255
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
256
  },
257
  {
258
  "id": "main",
259
  "name": "MatMulNBitsMlp.GateUp",
260
  "shader": "mlp-gate-up.wgsl.jinja",
261
+ "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
262
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
263
+ "subgroupCollectivesWidth": "portable"
264
  }
265
  ]
266
  },
267
  {
268
  "id": "staged_norm_nogb_ub",
269
  "priority": 10,
270
+ "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "not present.skipT", "not present.residualT"],
271
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
272
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
273
  "passes": [
274
  {
275
  "id": "norm",
276
  "name": "MatMulNBitsMlp.RmsNorm",
277
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
278
+ "bindings": ["a", "norm_scale", "normed", "params"],
279
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
280
  },
281
  {
282
  "id": "main",
283
  "name": "MatMulNBitsMlp.GateUp",
284
  "shader": "mlp-gate-up.wgsl.jinja",
285
+ "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
286
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
287
+ "subgroupCollectivesWidth": "portable"
288
  }
289
  ]
290
  },
291
  {
292
  "id": "staged_skip_nogb_ub",
293
  "priority": 10,
294
+ "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "present.skipT", "not present.residualT"],
295
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
296
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
297
  "passes": [
298
  {
299
  "id": "norm",
300
  "name": "MatMulNBitsMlp.RmsNorm",
301
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
302
+ "bindings": ["a", "skip", "norm_scale", "normed", "params"],
303
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
304
  },
305
  {
306
  "id": "main",
307
  "name": "MatMulNBitsMlp.GateUp",
308
  "shader": "mlp-gate-up.wgsl.jinja",
309
+ "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
310
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
311
+ "subgroupCollectivesWidth": "portable"
312
  }
313
  ]
314
  },
315
  {
316
  "id": "staged_skipsum_nogb_ub",
317
  "priority": 10,
318
+ "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "present.skipT", "present.residualT"],
319
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
320
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
321
  "passes": [
322
  {
323
  "id": "norm",
324
  "name": "MatMulNBitsMlp.RmsNorm",
325
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
326
+ "bindings": ["a", "skip", "norm_scale", "normed", "residual", "params"],
327
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
328
  },
329
  {
330
  "id": "main",
331
  "name": "MatMulNBitsMlp.GateUp",
332
  "shader": "mlp-gate-up.wgsl.jinja",
333
+ "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
334
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
335
+ "subgroupCollectivesWidth": "portable"
336
  }
337
  ]
338
  },
339
  {
340
  "id": "staged_norm_gb_noub",
341
  "priority": 10,
342
+ "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "not present.skipT", "not present.residualT"],
343
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
344
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
345
  "passes": [
346
  {
347
  "id": "norm",
348
  "name": "MatMulNBitsMlp.RmsNorm",
349
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
350
+ "bindings": ["a", "norm_scale", "normed", "params"],
351
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
352
  },
353
  {
354
  "id": "main",
355
  "name": "MatMulNBitsMlp.GateUp",
356
  "shader": "mlp-gate-up.wgsl.jinja",
357
+ "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
358
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
359
+ "subgroupCollectivesWidth": "portable"
360
  }
361
  ]
362
  },
363
  {
364
  "id": "staged_skip_gb_noub",
365
  "priority": 10,
366
+ "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "present.skipT", "not present.residualT"],
367
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
368
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
369
  "passes": [
370
  {
371
  "id": "norm",
372
  "name": "MatMulNBitsMlp.RmsNorm",
373
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
374
+ "bindings": ["a", "skip", "norm_scale", "normed", "params"],
375
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
376
  },
377
  {
378
  "id": "main",
379
  "name": "MatMulNBitsMlp.GateUp",
380
  "shader": "mlp-gate-up.wgsl.jinja",
381
+ "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
382
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
383
+ "subgroupCollectivesWidth": "portable"
384
  }
385
  ]
386
  },
387
  {
388
  "id": "staged_skipsum_gb_noub",
389
  "priority": 10,
390
+ "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "present.skipT", "present.residualT"],
391
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
392
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
393
  "passes": [
394
  {
395
  "id": "norm",
396
  "name": "MatMulNBitsMlp.RmsNorm",
397
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
398
+ "bindings": ["a", "skip", "norm_scale", "normed", "residual", "params"],
399
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
400
  },
401
  {
402
  "id": "main",
403
  "name": "MatMulNBitsMlp.GateUp",
404
  "shader": "mlp-gate-up.wgsl.jinja",
405
+ "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
406
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
407
+ "subgroupCollectivesWidth": "portable"
408
  }
409
  ]
410
  },
411
  {
412
  "id": "staged_norm_gb_ub",
413
  "priority": 10,
414
+ "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "not present.skipT", "not present.residualT"],
415
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
416
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
417
  "passes": [
418
  {
419
  "id": "norm",
420
  "name": "MatMulNBitsMlp.RmsNorm",
421
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
422
+ "bindings": ["a", "norm_scale", "normed", "params"],
423
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
424
  },
425
  {
426
  "id": "main",
427
  "name": "MatMulNBitsMlp.GateUp",
428
  "shader": "mlp-gate-up.wgsl.jinja",
429
+ "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
430
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
431
+ "subgroupCollectivesWidth": "portable"
432
  }
433
  ]
434
  },
435
  {
436
  "id": "staged_skip_gb_ub",
437
  "priority": 10,
438
+ "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "present.skipT", "not present.residualT"],
439
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
440
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
441
  "passes": [
442
  {
443
  "id": "norm",
444
  "name": "MatMulNBitsMlp.RmsNorm",
445
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
446
+ "bindings": ["a", "skip", "norm_scale", "normed", "params"],
447
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
448
  },
449
  {
450
  "id": "main",
451
  "name": "MatMulNBitsMlp.GateUp",
452
  "shader": "mlp-gate-up.wgsl.jinja",
453
+ "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
454
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
455
+ "subgroupCollectivesWidth": "portable"
456
  }
457
  ]
458
  },
459
  {
460
  "id": "staged_skipsum_gb_ub",
461
  "priority": 10,
462
+ "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "present.skipT", "present.residualT"],
463
+ "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
464
  "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
465
  "passes": [
466
  {
467
  "id": "norm",
468
  "name": "MatMulNBitsMlp.RmsNorm",
469
  "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
470
+ "bindings": ["a", "skip", "norm_scale", "normed", "residual", "params"],
471
+ "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
472
  },
473
  {
474
  "id": "main",
475
  "name": "MatMulNBitsMlp.GateUp",
476
  "shader": "mlp-gate-up.wgsl.jinja",
477
+ "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
478
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
479
+ "subgroupCollectivesWidth": "portable"
480
  }
481
  ]
482
  },
483
  {
484
  "id": "fused_norm_nogb_noub",
485
  "priority": 30,
486
+ "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
487
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 7 } },
488
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
489
  "passes": [
490
  {
491
  "id": "main",
492
  "name": "MatMulNBitsMlp.FusedDecode",
493
  "shader": "mlp-gate-up.wgsl.jinja",
494
+ "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
495
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
496
+ "subgroupCollectivesWidth": "portable"
497
  }
498
  ]
499
  },
500
  {
501
  "id": "fused_skip_nogb_noub",
502
  "priority": 30,
503
+ "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
504
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
505
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
506
  "passes": [
507
  {
508
  "id": "main",
509
  "name": "MatMulNBitsMlp.FusedDecode",
510
  "shader": "mlp-gate-up.wgsl.jinja",
511
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
512
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
513
+ "subgroupCollectivesWidth": "portable"
514
  }
515
  ]
516
  },
517
  {
518
  "id": "fused_skipsum_nogb_noub",
519
  "priority": 30,
520
+ "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "present.skipT", "present.residualT"],
521
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
522
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
523
  "passes": [
524
  {
525
  "id": "main",
526
  "name": "MatMulNBitsMlp.FusedDecode",
527
  "shader": "mlp-gate-up.wgsl.jinja",
528
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y", "residual"],
529
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
530
+ "subgroupCollectivesWidth": "portable"
531
  }
532
  ]
533
  },
534
  {
535
  "id": "fused_norm_nogb_ub",
536
  "priority": 30,
537
+ "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
538
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
539
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
540
  "passes": [
541
  {
542
  "id": "main",
543
  "name": "MatMulNBitsMlp.FusedDecode",
544
  "shader": "mlp-gate-up.wgsl.jinja",
545
+ "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
546
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
547
+ "subgroupCollectivesWidth": "portable"
548
  }
549
  ]
550
  },
551
  {
552
  "id": "fused_skip_nogb_ub",
553
  "priority": 30,
554
+ "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
555
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
556
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
557
  "passes": [
558
  {
559
  "id": "main",
560
  "name": "MatMulNBitsMlp.FusedDecode",
561
  "shader": "mlp-gate-up.wgsl.jinja",
562
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
563
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
564
+ "subgroupCollectivesWidth": "portable"
565
  }
566
  ]
567
  },
568
  {
569
  "id": "fused_skipsum_nogb_ub",
570
  "priority": 30,
571
+ "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "present.skipT", "present.residualT"],
572
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
573
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
574
  "passes": [
575
  {
576
  "id": "main",
577
  "name": "MatMulNBitsMlp.FusedDecode",
578
  "shader": "mlp-gate-up.wgsl.jinja",
579
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y", "residual"],
580
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
581
+ "subgroupCollectivesWidth": "portable"
582
  }
583
  ]
584
  },
585
  {
586
  "id": "fused_norm_gb_noub",
587
  "priority": 30,
588
+ "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
589
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
590
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
591
  "passes": [
592
  {
593
  "id": "main",
594
  "name": "MatMulNBitsMlp.FusedDecode",
595
  "shader": "mlp-gate-up.wgsl.jinja",
596
+ "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
597
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
598
+ "subgroupCollectivesWidth": "portable"
599
  }
600
  ]
601
  },
602
  {
603
  "id": "fused_skip_gb_noub",
604
  "priority": 30,
605
+ "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
606
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
607
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
608
  "passes": [
609
  {
610
  "id": "main",
611
  "name": "MatMulNBitsMlp.FusedDecode",
612
  "shader": "mlp-gate-up.wgsl.jinja",
613
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
614
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
615
+ "subgroupCollectivesWidth": "portable"
616
  }
617
  ]
618
  },
619
  {
620
  "id": "fused_skipsum_gb_noub",
621
  "priority": 30,
622
+ "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "present.skipT", "present.residualT"],
623
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
624
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
625
  "passes": [
626
  {
627
  "id": "main",
628
  "name": "MatMulNBitsMlp.FusedDecode",
629
  "shader": "mlp-gate-up.wgsl.jinja",
630
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y", "residual"],
631
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
632
+ "subgroupCollectivesWidth": "portable"
633
  }
634
  ]
635
  },
636
  {
637
  "id": "fused_norm_gb_ub",
638
  "priority": 30,
639
+ "when": ["normContractOk", "biasPresence_gb_ub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
640
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
641
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
642
  "passes": [
643
  {
644
  "id": "main",
645
  "name": "MatMulNBitsMlp.FusedDecode",
646
  "shader": "mlp-gate-up.wgsl.jinja",
647
+ "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
648
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
649
+ "subgroupCollectivesWidth": "portable"
650
  }
651
  ]
652
  },
653
  {
654
  "id": "fused_skip_gb_ub",
655
  "priority": 30,
656
+ "when": ["normContractOk", "biasPresence_gb_ub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
657
  "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
658
+ "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
 
659
  "passes": [
660
  {
661
  "id": "main",
662
  "name": "MatMulNBitsMlp.FusedDecode",
663
  "shader": "mlp-gate-up.wgsl.jinja",
664
+ "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
665
+ "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
666
+ "subgroupCollectivesWidth": "portable"
667
  }
668
  ]
669
  }
build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja CHANGED
@@ -1,14 +1,10 @@
1
- {% if usesF16 %}
2
- enable f16;
3
- {% endif %}
4
  {{ env.wgsl.resourceDeclarations }}
5
 
6
- // Shared MatMulNBits MLP/QKV normalization pass.
7
  // normed[row, d] = (A + skip)[row, d] * inverseSqrt(mean_d((A + skip)^2) + eps) * norm_scale[d]
8
  // One workgroup owns one row. Every intermediate stays in f32 and `normed` is an
9
- // f32 scratch tensor, so the projection pass reads exactly what the fused
10
- // single-dispatch kernel keeps in registers -- the two paths agree for float16
11
- // inputs instead of differing by one narrowing.
12
  const HIDDEN: u32 = {{ hidden }}u;
13
  const WG: u32 = {{ workgroupSize }}u;
14
  const EPSILON: f32 = {{ epsilon }};
@@ -88,11 +84,10 @@ fn row_value(index: u32) -> f32 {
88
 
89
  @compute @workgroup_size(WG, 1, 1)
90
  fn main(@builtin(workgroup_id) wg: vec3<u32>,
91
- @builtin(num_workgroups) nwg: vec3<u32>,
92
  @builtin(local_invocation_id) lid: vec3<u32>) {
93
  // 2D-folded row index: wg.y carries the high bits past the
94
- // maxComputeWorkgroupsPerDimension dispatch limit. Reduces to wg.x when nwg.y == 1.
95
- let row = wg.x + wg.y * nwg.x;
96
  if (row >= params.rows) {
97
  return;
98
  }
 
 
 
 
1
  {{ env.wgsl.resourceDeclarations }}
2
 
3
+ // Fused RMS-normalization pass.
4
  // normed[row, d] = (A + skip)[row, d] * inverseSqrt(mean_d((A + skip)^2) + eps) * norm_scale[d]
5
  // One workgroup owns one row. Every intermediate stays in f32 and `normed` is an
6
+ // f32 scratch tensor. The following projection therefore consumes the
7
+ // normalized values without an intervening storage-type narrowing.
 
8
  const HIDDEN: u32 = {{ hidden }}u;
9
  const WG: u32 = {{ workgroupSize }}u;
10
  const EPSILON: f32 = {{ epsilon }};
 
84
 
85
  @compute @workgroup_size(WG, 1, 1)
86
  fn main(@builtin(workgroup_id) wg: vec3<u32>,
 
87
  @builtin(local_invocation_id) lid: vec3<u32>) {
88
  // 2D-folded row index: wg.y carries the high bits past the
89
+ // per-axis dispatch fold width. Reduces to wg.x when the dispatch does not fold.
90
+ let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
91
  if (row >= params.rows) {
92
  return;
93
  }
build/webgpu/metadata.json CHANGED
@@ -1,19 +1,50 @@
1
  {
2
  "name": "com.microsoft.MatMulNBitsMlp",
3
- "id": "_com_microsoft_matmulnbitsmlp_webgpu_5f0af61",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "FmfZaHEP/sBLMvGjNKkLAHUkkW9+LWNCTLnfAniZXZY=",
11
- "manifest.json": "2PUEf3quWni+CsZC4X4XMrR/GVu+mhWWWPqOjzAbbyU=",
12
- "matmul-nbits-fused-rms-norm.wgsl.jinja": "4TI9Mc/RZWVgMl6NVG2iwe5OpvqQmswhe6xtGvctNmI=",
13
- "mlp-gate-up.wgsl.jinja": "Zncc9Fsa1Vja8Dz8384/+MSAlJXnpb8OwIoGO1ht0rs=",
14
- "test.json": "sKKYxHcXolfGhck7PQNdGq3iSM8B58IYf/w5hl6JPLg="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
18
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MatMulNBitsMlp" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  }
 
1
  {
2
  "name": "com.microsoft.MatMulNBitsMlp",
3
+ "id": "_com_microsoft_matmulnbitsmlp_webgpu_d711f87",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "teEo4C2oaUqxuIKr6wLZW6vFlmI9gTl/akkWJhjWxqg=",
11
+ "manifest.json": "ceiKDb7Itix3E2Ab3NY1WqE4vb5P3kKtrPxTrImqQQk=",
12
+ "matmul-nbits-fused-rms-norm.wgsl.jinja": "4lOdB+RprQh3iv29i6RV8UWkn8S5y1aK6Te5nJpxuEk=",
13
+ "mlp-gate-up.wgsl.jinja": "VuwljvMtS5vV09dkhewYJLe6s6Syr6zEtOTiGQMMukI=",
14
+ "test.json": "x1T4v2UMpwHHuiJcAq6Y8p4TkcyvZwz5pbpbV38m5D4="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
+ "webgpu": {
19
+ "manifestSpec": "2.0",
20
+ "variants": {
21
+ "plain_nogb_noub": ["mlp-gate-up.wgsl.jinja"],
22
+ "plain_nogb_ub": ["mlp-gate-up.wgsl.jinja"],
23
+ "plain_gb_noub": ["mlp-gate-up.wgsl.jinja"],
24
+ "plain_gb_ub": ["mlp-gate-up.wgsl.jinja"],
25
+ "staged_norm_nogb_noub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
26
+ "staged_skip_nogb_noub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
27
+ "staged_skipsum_nogb_noub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
28
+ "staged_norm_nogb_ub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
29
+ "staged_skip_nogb_ub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
30
+ "staged_skipsum_nogb_ub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
31
+ "staged_norm_gb_noub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
32
+ "staged_skip_gb_noub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
33
+ "staged_skipsum_gb_noub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
34
+ "staged_norm_gb_ub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
35
+ "staged_skip_gb_ub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
36
+ "staged_skipsum_gb_ub": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "mlp-gate-up.wgsl.jinja"],
37
+ "fused_norm_nogb_noub": ["mlp-gate-up.wgsl.jinja"],
38
+ "fused_skip_nogb_noub": ["mlp-gate-up.wgsl.jinja"],
39
+ "fused_skipsum_nogb_noub": ["mlp-gate-up.wgsl.jinja"],
40
+ "fused_norm_nogb_ub": ["mlp-gate-up.wgsl.jinja"],
41
+ "fused_skip_nogb_ub": ["mlp-gate-up.wgsl.jinja"],
42
+ "fused_skipsum_nogb_ub": ["mlp-gate-up.wgsl.jinja"],
43
+ "fused_norm_gb_noub": ["mlp-gate-up.wgsl.jinja"],
44
+ "fused_skip_gb_noub": ["mlp-gate-up.wgsl.jinja"],
45
+ "fused_skipsum_gb_noub": ["mlp-gate-up.wgsl.jinja"],
46
+ "fused_norm_gb_ub": ["mlp-gate-up.wgsl.jinja"],
47
+ "fused_skip_gb_ub": ["mlp-gate-up.wgsl.jinja"]
48
+ }
49
+ }
50
  }
build/webgpu/mlp-gate-up.wgsl.jinja CHANGED
@@ -1,24 +1,22 @@
1
  {% macro matmul_nbits_packed_code(fn="packed_weight", buffer="b", kBlocks="params.kBlocks", blobSize="params.blobSize", bits=4) %}
2
  fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
3
  {% if bits == 2 %}
4
- let byte_index = offset / 4u;
5
- let shift = (offset % 4u) * 2u;
6
- let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + byte_index;
7
- return ({{ buffer }}[packed_index] >> shift) & 3u;
8
  {% elif bits == 4 %}
9
- let byte_index = offset / 2u;
10
- let shift = (offset % 2u) * 4u;
11
- let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + byte_index;
12
- return ({{ buffer }}[packed_index] >> shift) & 15u;
13
  {% else %}
14
- let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
15
- return {{ buffer }}[packed_index] & 255u;
16
  {% endif %}
17
  }
18
  {%- endmacro %}
19
 
20
- {% if usesF16 %}
21
- enable f16;
22
  {% endif %}
23
  {{ env.wgsl.resourceDeclarations }}
24
 
@@ -36,45 +34,84 @@ enable f16;
36
  // across their accumulators.
37
  const K: u32 = {{ K }}u;
38
  const N: u32 = {{ N }}u;
 
 
 
39
  const BLOCK_SIZE: u32 = {{ blockSize }}u;
 
40
  const KBLOCKS: u32 = {{ kBlocks }}u;
41
  const BLOB_SIZE: u32 = {{ blobSize }}u;
 
 
 
 
 
 
 
42
  const TILE_N: u32 = {{ tileN }}u;
43
  const LANES: u32 = {{ lanes }}u;
44
  const ROW_TILE: u32 = {{ rowTile }}u;
45
- const ROWS: u32 = {{ rows }}u;
46
  const WG: u32 = TILE_N * LANES;
47
  const ZERO: f32 = {{ defaultZero }};
48
- // Codes per logical byte and the mask for one code. Physical uint8 storage uses
49
- // a u32 slot, so packing increases the number of codes returned by each load.
 
 
 
50
  const BITS: u32 = {{ bits }}u;
 
51
  const CODES_PER_BYTE: u32 = {{ codesPerByte }}u;
52
  const CODE_MASK: u32 = {{ codeMask }}u;
 
 
 
 
 
 
 
 
 
 
53
  {% if inlineNorm %}
54
  const EPSILON: f32 = {{ epsilon }};
55
  {% endif %}
56
 
57
  {% for stream in ["gate", "up"] %}
 
58
  {{ matmul_nbits_packed_code(fn=stream ~ "_code", buffer=stream ~ "_b", kBlocks="KBLOCKS", blobSize="BLOB_SIZE", bits=bits) }}
 
 
59
  // Decode two consecutive reduction-axis codes. Below 8 bits an even offset and
60
- // its successor share one stored byte; at 8 bits they occupy adjacent slots. An
61
- // odd offset would straddle bytes, so callers advance by two from an even start.
 
 
62
  fn {{ stream }}_code_pair(n: u32, block: u32, offset: u32) -> vec2<u32> {
63
- let base = (n * KBLOCKS + block) * BLOB_SIZE;
64
- let shift = (offset % CODES_PER_BYTE) * BITS;
65
- let lo = {{ stream }}_b[base + offset / CODES_PER_BYTE];
66
- let hi = {{ "lo" if codesPerByte > 1 else (stream ~ "_b[base + offset + 1u]") }};
67
- return vec2<u32>((lo >> shift) & CODE_MASK,
68
- (hi >> {{ "(shift + BITS)" if codesPerByte > 1 else "0u" }}) & CODE_MASK);
69
  }
 
70
  {% endfor %}
71
 
72
  fn silu(x: f32) -> f32 {
73
  return x / (1.0 + exp(-x));
74
  }
75
 
 
 
 
 
 
 
 
 
 
76
  var<workgroup> red_gate: array<f32, WG * ROW_TILE>;
77
  var<workgroup> red_up: array<f32, WG * ROW_TILE>;
 
78
  {% if inlineNorm %}
79
  var<workgroup> partial: array<f32, WG>;
80
  var<workgroup> row_inv: f32;
@@ -153,6 +190,255 @@ fn row_value(index: u32) -> f32 {
153
 
154
  {% macro act(b, k) %}{% if inlineNorm %}row_value({{ b }} + {{ k }}) * row_inv * f32(norm_scale[{{ k }}]){% elif fromNormed %}normed[{{ b }} + {{ k }}]{% else %}f32(a[{{ b }} + {{ k }}]){% endif %}{%- endmacro %}
155
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
156
  {% macro walk_block(guarded) %}
157
  for (var offset = lane * 2u; offset + 1u < BLOCK_SIZE; offset = offset + LANES * 2u) {
158
  let k = k_base + offset;
@@ -202,29 +488,6 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
202
  let base_{{ r }} = min(row0 + {{ r }}u, ROWS - 1u) * K;
203
  {% endfor %}
204
 
205
- {% if inlineNorm %}
206
- var local_sq = 0.0;
207
- for (var d = tid; d < K; d = d + WG) {
208
- let value = row_value(base_0 + d);
209
- local_sq = local_sq + value * value;
210
- }
211
- let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(K) + EPSILON);
212
- if (tid == 0u) {
213
- row_inv = inv;
214
- }
215
- // Separates the reduction's readers of partial[0] from the projection's
216
- // reuse of the same workgroup array below.
217
- workgroupBarrier();
218
- {% if writeResidual %}
219
- // Every N tile computes the same residual row; only the first one stores it,
220
- // so the tiles never write the same location.
221
- if (wg.x == 0u) {
222
- for (var d = tid; d < K; d = d + WG) {
223
- residual[base_0 + d] = {{ scalar }}(row_value(base_0 + d));
224
- }
225
- }
226
- {% endif %}
227
- {% endif %}
228
 
229
  {% for r in range(rowTile) %}
230
  var acc_gate_{{ r }} = 0.0;
@@ -287,4 +550,4 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
287
  }
288
  {% endfor %}
289
  }
290
- }
 
1
  {% macro matmul_nbits_packed_code(fn="packed_weight", buffer="b", kBlocks="params.kBlocks", blobSize="params.blobSize", bits=4) %}
2
  fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
3
  {% if bits == 2 %}
4
+ let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + (offset >> 2u);
5
+ let shift = (byte_index & 3u) * 8u + (offset & 3u) * 2u;
6
+ return ({{ buffer }}[byte_index >> 2u] >> shift) & 3u;
 
7
  {% elif bits == 4 %}
8
+ let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + (offset >> 1u);
9
+ let shift = (byte_index & 3u) * 8u + (offset & 1u) * 4u;
10
+ return ({{ buffer }}[byte_index >> 2u] >> shift) & 15u;
 
11
  {% else %}
12
+ let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
13
+ return ({{ buffer }}[byte_index >> 2u] >> ((byte_index & 3u) * 8u)) & 255u;
14
  {% endif %}
15
  }
16
  {%- endmacro %}
17
 
18
+ {% if rowTile == 1 and useSubgroups %}
19
+ enable subgroups;
20
  {% endif %}
21
  {{ env.wgsl.resourceDeclarations }}
22
 
 
34
  // across their accumulators.
35
  const K: u32 = {{ K }}u;
36
  const N: u32 = {{ N }}u;
37
+ {% if not decodeVec %}
38
+ // The pair and tile walks address codes through the block; the vector walk
39
+ // addresses whole vector groups and never needs it.
40
  const BLOCK_SIZE: u32 = {{ blockSize }}u;
41
+ {% endif %}
42
  const KBLOCKS: u32 = {{ kBlocks }}u;
43
  const BLOB_SIZE: u32 = {{ blobSize }}u;
44
+ {% set gemvWalk = rowTile == 1 %}
45
+ {% if gemvWalk %}
46
+ const N_COLS: u32 = {{ decodeNCols }}u;
47
+ const WG: u32 = {{ decodeWorkgroupSize }}u;
48
+ const ROWS: u32 = {{ rowCount }}u;
49
+ const ZERO: f32 = {{ defaultZero }};
50
+ {% else %}
51
  const TILE_N: u32 = {{ tileN }}u;
52
  const LANES: u32 = {{ lanes }}u;
53
  const ROW_TILE: u32 = {{ rowTile }}u;
54
+ const ROWS: u32 = {{ rowCount }}u;
55
  const WG: u32 = TILE_N * LANES;
56
  const ZERO: f32 = {{ defaultZero }};
57
+ {% endif %}
58
+ // Codes per blob byte and the mask for one code. The blob is bound in the
59
+ // packed storage layout -- four blob bytes per u32 word -- so a byte index is
60
+ // a word index plus a shift.
61
+ {% if not decodeVec %}
62
  const BITS: u32 = {{ bits }}u;
63
+ {% endif %}
64
  const CODES_PER_BYTE: u32 = {{ codesPerByte }}u;
65
  const CODE_MASK: u32 = {{ codeMask }}u;
66
+ {% if decodeVec %}
67
+ // One vec4<u32> of the packed blob is sixteen stored bytes, so it carries
68
+ // 16 * CODES_PER_BYTE codes -- and a whole quantization block is exactly
69
+ // BLOB_SIZE / 16 of them (this arm requires blobSize % 16 == 0). A vector
70
+ // therefore never straddles two blocks and the flat vector index needs no
71
+ // division: (n * KBLOCKS + block) * (BLOB_SIZE / 16) + slot == n * VEC_GROUPS + group.
72
+ const VEC_PER_BLOCK: u32 = BLOB_SIZE / 16u;
73
+ const VEC_GROUPS: u32 = KBLOCKS * VEC_PER_BLOCK;
74
+ const CODES_PER_VEC: u32 = 16u * CODES_PER_BYTE;
75
+ {% endif %}
76
  {% if inlineNorm %}
77
  const EPSILON: f32 = {{ epsilon }};
78
  {% endif %}
79
 
80
  {% for stream in ["gate", "up"] %}
81
+ {% if not gemvWalk %}
82
  {{ matmul_nbits_packed_code(fn=stream ~ "_code", buffer=stream ~ "_b", kBlocks="KBLOCKS", blobSize="BLOB_SIZE", bits=bits) }}
83
+ {% endif %}
84
+ {% if not decodeVec %}
85
  // Decode two consecutive reduction-axis codes. Below 8 bits an even offset and
86
+ // its successor share one stored byte; at 8 bits they occupy adjacent bytes of
87
+ // one word, because an even byte index in a 4-byte-aligned blob never ends a
88
+ // word. An odd offset would straddle bytes, so callers advance by two from an
89
+ // even start.
90
  fn {{ stream }}_code_pair(n: u32, block: u32, offset: u32) -> vec2<u32> {
91
+ let byte_index = (n * KBLOCKS + block) * BLOB_SIZE + offset / CODES_PER_BYTE;
92
+ let shift = (byte_index & 3u) * 8u + (offset % CODES_PER_BYTE) * BITS;
93
+ let word = {{ stream }}_b[byte_index >> 2u];
94
+ return vec2<u32>((word >> shift) & CODE_MASK, (word >> (shift + BITS)) & CODE_MASK);
 
 
95
  }
96
+ {% endif %}
97
  {% endfor %}
98
 
99
  fn silu(x: f32) -> f32 {
100
  return x / (1.0 + exp(-x));
101
  }
102
 
103
+ {% if gemvWalk %}
104
+ {% set colGroups = 2 if decodeNCols == 8 else 1 %}
105
+ {% set comps = ["x", "y", "z", "w"] %}
106
+ {% for gi in range(colGroups) %}
107
+ {% set sfx = "" if gi == 0 else gi %}
108
+ var<workgroup> red_gate{{ sfx }}: array<vec4<f32>, WG>;
109
+ var<workgroup> red_up{{ sfx }}: array<vec4<f32>, WG>;
110
+ {% endfor %}
111
+ {% else %}
112
  var<workgroup> red_gate: array<f32, WG * ROW_TILE>;
113
  var<workgroup> red_up: array<f32, WG * ROW_TILE>;
114
+ {% endif %}
115
  {% if inlineNorm %}
116
  var<workgroup> partial: array<f32, WG>;
117
  var<workgroup> row_inv: f32;
 
190
 
191
  {% macro act(b, k) %}{% if inlineNorm %}row_value({{ b }} + {{ k }}) * row_inv * f32(norm_scale[{{ k }}]){% elif fromNormed %}normed[{{ b }} + {{ k }}]{% else %}f32(a[{{ b }} + {{ k }}]){% endif %}{%- endmacro %}
192
 
193
+ {% if gemvWalk %}
194
+ @compute @workgroup_size(WG, 1, 1)
195
+ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
196
+ @builtin(subgroup_invocation_id) sgLane: u32, @builtin(subgroup_size) sgSize: u32{% endif %}) {
197
+ let tid = lid.x;
198
+ // Rows past the end of the batch clamp onto the last real row; the store
199
+ // guard drops them.
200
+ let row = min(wg.y, ROWS - 1u);
201
+ let base_0 = row * K;
202
+ let col_base = wg.x * N_COLS;
203
+ {% if inlineNorm %}
204
+ var local_sq = 0.0;
205
+ for (var d = tid; d < K; d = d + WG) {
206
+ let value = row_value(base_0 + d);
207
+ local_sq = local_sq + value * value;
208
+ }
209
+ let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(K) + EPSILON);
210
+ if (tid == 0u) {
211
+ row_inv = inv;
212
+ }
213
+ // Separates the reduction's readers of partial[0] from the projection's
214
+ // reuse of the same workgroup array below.
215
+ workgroupBarrier();
216
+ {% if writeResidual %}
217
+ // Every N tile computes the same residual row; only the first one stores it,
218
+ // so the tiles never write the same location.
219
+ if (wg.x == 0u) {
220
+ for (var d = tid; d < K; d = d + WG) {
221
+ residual[base_0 + d] = {{ scalar }}(row_value(base_0 + d));
222
+ }
223
+ }
224
+ {% endif %}
225
+ {% endif %}
226
+ // Whole workgroups past the last column return here, after the normalization
227
+ // barriers above and before the reduction's.
228
+ if (col_base >= N) {
229
+ return;
230
+ }
231
+ {% for gi in range(colGroups) %}
232
+ {% set sfx = "" if gi == 0 else gi %}
233
+ var acc_gate{{ sfx }} = vec4<f32>(0.0);
234
+ var acc_up{{ sfx }} = vec4<f32>(0.0);
235
+ {% endfor %}
236
+ {% if decodeVec %}
237
+ {% set codesPerWord = 4 * codesPerByte %}
238
+ {% if laneGroups == 2 %}
239
+ // Two lane halves traverse the same vector groups: the lower half accumulates
240
+ // the gate projection and the upper half accumulates the up projection. The
241
+ // inactive projection accumulator remains zero and folds harmlessly.
242
+ const K_LANES: u32 = WG / 2u;
243
+ let laneProj = tid / K_LANES;
244
+ for (var g = tid % K_LANES; g < VEC_GROUPS; g = g + K_LANES) {
245
+ {% else %}
246
+ for (var g = tid; g < VEC_GROUPS; g = g + WG) {
247
+ {% endif %}
248
+ let block = g / VEC_PER_BLOCK;
249
+ {% if actVec4 and fromNormed %}
250
+ // K is a whole number of blocks here, so the group lands entirely inside the
251
+ // row and the staged activations come from CODES_PER_VEC / 4 vector loads.
252
+ let vbase = base_0 / 4u + g * (CODES_PER_VEC / 4u);
253
+ {% for v in range(codesPerWord) %}
254
+ let av{{ v }} = normed[vbase + {{ v }}u];
255
+ {% endfor %}
256
+ {% for v in range(codesPerWord) %}
257
+ {% for c in range(4) %}
258
+ let a{{ v * 4 + c }} = av{{ v }}.{{ ["x", "y", "z", "w"][c] }};
259
+ {% endfor %}
260
+ {% endfor %}
261
+ {% else %}
262
+ let k0 = g * CODES_PER_VEC;
263
+ {% for j in range(16 * codesPerByte) %}
264
+ // A trailing partial block reaches past K; those codes contribute zero.
265
+ let a{{ j }} = select(0.0, {{ act("base_0", "min(k0 + " ~ j ~ "u, K - 1u)") }}, k0 + {{ j }}u < K);
266
+ {% endfor %}
267
+ {% endif %}
268
+ let asum = {% for j in range(16 * codesPerByte) %}{{ " + " if j > 0 else "" }}a{{ j }}{% endfor %};
269
+ {% for gi in range(colGroups) %}
270
+ {% set sfx = "" if gi == 0 else gi %}
271
+ {% for c in range(4) %}
272
+ {% set i = gi * 4 + c %}
273
+ {% set comp = comps[c] %}
274
+ {% if i == 0 %}
275
+ {
276
+ {% else %}
277
+ if (col_base + {{ i }}u < N) {
278
+ {% endif %}
279
+ let n = col_base + {{ i }}u;
280
+ {% if laneGroups == 2 %}
281
+ // The lane's half selects the projection; the branch is uniform across
282
+ // each half and the loads it skips are the other half's.
283
+ if (laneProj == 0u) {
284
+ let gate_words = gate_b[n * VEC_GROUPS + g];
285
+ let gate_scale = f32(gate_scales[n * KBLOCKS + block]);
286
+ var gate_dot = 0.0;
287
+ {% for w in range(4) %}
288
+ {% for h in range(codesPerWord) %}
289
+ gate_dot = gate_dot + a{{ w * codesPerWord + h }} * f32((gate_words.{{ comps[w] }} >> {{ h * bits }}u) & CODE_MASK);
290
+ {% endfor %}
291
+ {% endfor %}
292
+ acc_gate{{ sfx }}.{{ comp }} = acc_gate{{ sfx }}.{{ comp }} + (gate_dot - ZERO * asum) * gate_scale;
293
+ } else {
294
+ let up_words = up_b[n * VEC_GROUPS + g];
295
+ let up_scale = f32(up_scales[n * KBLOCKS + block]);
296
+ var up_dot = 0.0;
297
+ {% for w in range(4) %}
298
+ {% for h in range(codesPerWord) %}
299
+ up_dot = up_dot + a{{ w * codesPerWord + h }} * f32((up_words.{{ comps[w] }} >> {{ h * bits }}u) & CODE_MASK);
300
+ {% endfor %}
301
+ {% endfor %}
302
+ acc_up{{ sfx }}.{{ comp }} = acc_up{{ sfx }}.{{ comp }} + (up_dot - ZERO * asum) * up_scale;
303
+ }
304
+ {% else %}
305
+ let gate_words = gate_b[n * VEC_GROUPS + g];
306
+ let up_words = up_b[n * VEC_GROUPS + g];
307
+ let gate_scale = f32(gate_scales[n * KBLOCKS + block]);
308
+ let up_scale = f32(up_scales[n * KBLOCKS + block]);
309
+ var gate_dot = 0.0;
310
+ var up_dot = 0.0;
311
+ {% for w in range(4) %}
312
+ {% for h in range(codesPerWord) %}
313
+ gate_dot = gate_dot + a{{ w * codesPerWord + h }} * f32((gate_words.{{ comps[w] }} >> {{ h * bits }}u) & CODE_MASK);
314
+ up_dot = up_dot + a{{ w * codesPerWord + h }} * f32((up_words.{{ comps[w] }} >> {{ h * bits }}u) & CODE_MASK);
315
+ {% endfor %}
316
+ {% endfor %}
317
+ acc_gate{{ sfx }}.{{ comp }} = acc_gate{{ sfx }}.{{ comp }} + (gate_dot - ZERO * asum) * gate_scale;
318
+ acc_up{{ sfx }}.{{ comp }} = acc_up{{ sfx }}.{{ comp }} + (up_dot - ZERO * asum) * up_scale;
319
+ {% endif %}
320
+ }
321
+ {% endfor %}
322
+ {% endfor %}
323
+ }
324
+ {% else %}
325
+ for (var k = tid * 2u; k < K; k = k + WG * 2u) {
326
+ let block = k / BLOCK_SIZE;
327
+ let offset = k - block * BLOCK_SIZE;
328
+ let v0 = {{ act("base_0", "k") }};
329
+ // K need not be even; a code past the end contributes zero.
330
+ let v1 = select(0.0, {{ act("base_0", "min(k + 1u, K - 1u)") }}, k + 1u < K);
331
+ {% for c in range(4) %}
332
+ {% set comp = ["x", "y", "z", "w"][c] %}
333
+ {% if c == 0 %}
334
+ {
335
+ {% else %}
336
+ if (col_base + {{ c }}u < N) {
337
+ {% endif %}
338
+ let n = col_base + {{ c }}u;
339
+ let gate_codes = gate_code_pair(n, block, offset);
340
+ let up_codes = up_code_pair(n, block, offset);
341
+ let gate_scale = f32(gate_scales[n * KBLOCKS + block]);
342
+ let up_scale = f32(up_scales[n * KBLOCKS + block]);
343
+ acc_gate.{{ comp }} = acc_gate.{{ comp }}
344
+ + (v0 * (f32(gate_codes.x) - ZERO) + v1 * (f32(gate_codes.y) - ZERO)) * gate_scale;
345
+ acc_up.{{ comp }} = acc_up.{{ comp }}
346
+ + (v0 * (f32(up_codes.x) - ZERO) + v1 * (f32(up_codes.y) - ZERO)) * up_scale;
347
+ }
348
+ {% endfor %}
349
+ }
350
+ {% endif %}
351
+
352
+ {% if useSubgroups %}
353
+ // Subgroup fold: one collective per projection, then the WG / subgroup-size
354
+ // per-subgroup partials fold once through workgroup memory. The lanes of a
355
+ // subgroup are contiguous in local_invocation_id, and every lane reaches this
356
+ // point (the early return above is workgroup-uniform), so the collectives run
357
+ // in uniform control flow.
358
+ {% for gi in range(colGroups) %}
359
+ {% set sfx = "" if gi == 0 else gi %}
360
+ let sgGate{{ sfx }} = subgroupAdd(acc_gate{{ sfx }});
361
+ let sgUp{{ sfx }} = subgroupAdd(acc_up{{ sfx }});
362
+ {% endfor %}
363
+ if (sgLane == 0u) {
364
+ {% for gi in range(colGroups) %}
365
+ {% set sfx = "" if gi == 0 else gi %}
366
+ red_gate{{ sfx }}[tid / sgSize] = sgGate{{ sfx }};
367
+ red_up{{ sfx }}[tid / sgSize] = sgUp{{ sfx }};
368
+ {% endfor %}
369
+ }
370
+ workgroupBarrier();
371
+
372
+ if (tid == 0u && wg.y < ROWS) {
373
+ let subgroupCount = WG / sgSize;
374
+ {% for gi in range(colGroups) %}
375
+ {% set sfx = "" if gi == 0 else gi %}
376
+ var gate_total{{ sfx }} = red_gate{{ sfx }}[0];
377
+ var up_total{{ sfx }} = red_up{{ sfx }}[0];
378
+ {% endfor %}
379
+ for (var i = 1u; i < subgroupCount; i = i + 1u) {
380
+ {% for gi in range(colGroups) %}
381
+ {% set sfx = "" if gi == 0 else gi %}
382
+ gate_total{{ sfx }} = gate_total{{ sfx }} + red_gate{{ sfx }}[i];
383
+ up_total{{ sfx }} = up_total{{ sfx }} + red_up{{ sfx }}[i];
384
+ {% endfor %}
385
+ }
386
+ {% else %}
387
+ {% for gi in range(colGroups) %}
388
+ {% set sfx = "" if gi == 0 else gi %}
389
+ red_gate{{ sfx }}[tid] = acc_gate{{ sfx }};
390
+ red_up{{ sfx }}[tid] = acc_up{{ sfx }};
391
+ {% endfor %}
392
+ workgroupBarrier();
393
+ var stride = WG / 2u;
394
+ loop {
395
+ if (stride == 0u) {
396
+ break;
397
+ }
398
+ if (tid < stride) {
399
+ {% for gi in range(colGroups) %}
400
+ {% set sfx = "" if gi == 0 else gi %}
401
+ red_gate{{ sfx }}[tid] = red_gate{{ sfx }}[tid] + red_gate{{ sfx }}[tid + stride];
402
+ red_up{{ sfx }}[tid] = red_up{{ sfx }}[tid] + red_up{{ sfx }}[tid + stride];
403
+ {% endfor %}
404
+ }
405
+ stride = stride / 2u;
406
+ workgroupBarrier();
407
+ }
408
+
409
+ if (tid == 0u && wg.y < ROWS) {
410
+ {% for gi in range(colGroups) %}
411
+ {% set sfx = "" if gi == 0 else gi %}
412
+ let gate_total{{ sfx }} = red_gate{{ sfx }}[0];
413
+ let up_total{{ sfx }} = red_up{{ sfx }}[0];
414
+ {% endfor %}
415
+ {% endif %}
416
+ {% for gi in range(colGroups) %}
417
+ {% set sfx = "" if gi == 0 else gi %}
418
+ {% for c in range(4) %}
419
+ {% set i = gi * 4 + c %}
420
+ {% set comp = comps[c] %}
421
+ {% if i == 0 %}
422
+ {
423
+ {% else %}
424
+ if (col_base + {{ i }}u < N) {
425
+ {% endif %}
426
+ let n = col_base + {{ i }}u;
427
+ var gate_value = gate_total{{ sfx }}.{{ comp }};
428
+ var up_value = up_total{{ sfx }}.{{ comp }};
429
+ {% if hasGateBias %}
430
+ gate_value = gate_value + f32(gate_bias[n]);
431
+ {% endif %}
432
+ {% if hasUpBias %}
433
+ up_value = up_value + f32(up_bias[n]);
434
+ {% endif %}
435
+ y[row * N + n] = {{ scalar }}(silu(gate_value) * up_value);
436
+ }
437
+ {% endfor %}
438
+ {% endfor %}
439
+ }
440
+ }
441
+ {% else %}
442
  {% macro walk_block(guarded) %}
443
  for (var offset = lane * 2u; offset + 1u < BLOCK_SIZE; offset = offset + LANES * 2u) {
444
  let k = k_base + offset;
 
488
  let base_{{ r }} = min(row0 + {{ r }}u, ROWS - 1u) * K;
489
  {% endfor %}
490
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
491
 
492
  {% for r in range(rowTile) %}
493
  var acc_gate_{{ r }} = 0.0;
 
550
  }
551
  {% endfor %}
552
  }
553
+ }{% endif %}
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.MatMulNBitsMlp",
3
  "fixtureArrays": {
4
  "pinned_plain_gb_ub_input_aT": [0.9181, 1.1011, 1.115, 0.9522, 0.6333, 0.2043, -0.2696, -0.7143, -1.0586, -1.2453, -1.2411, -1.0421, -0.6753, -0.1947, 0.3263, 0.8068, 1.1707, 1.3587, 1.3385, 1.1098, 0.7051, 0.1852, -0.3707, -0.8762, -1.2521, -1.4384, -1.4042, -1.1526, -0.7208, -0.1744, 0.4029, 0.9218],
5
  "pinned_plain_gb_ub_input_gateBT": [52, 93, 210, 163, 160, 89, 30, 255, 204, 21, 42, 27, 184, 145, 246, 247, 100, 205, 130, 147, 208, 201, 206, 239, 252, 133, 218, 11, 232, 1, 166, 231],
@@ -1230,7 +1229,9 @@
1230
  },
1231
  {
1232
  "name": "block32_decode",
1233
- "provenance": { "notes": "block_size 32, the size ONNX Runtime's fused decode kernel is specialized for." },
 
 
1234
  "attrs": { "K": 64, "N": 8, "bits": 4, "block_size": 32, "activation": "silu" },
1235
  "inputs": {
1236
  "aT": {
@@ -1431,7 +1432,7 @@
1431
  {
1432
  "name": "pinned_plain_gb_ub",
1433
  "provenance": {
1434
- "notes": "Expected values computed by an independent implementation written from the ONNX Runtime schema text alone, so this case checks the trusted reference as well as the kernels. No normalization: the projections read A directly."
1435
  },
1436
  "attrs": { "K": 16, "N": 4, "bits": 4, "block_size": 8, "activation": "silu" },
1437
  "inputs": {
@@ -1487,7 +1488,7 @@
1487
  {
1488
  "name": "pinned_norm_nogb_noub",
1489
  "provenance": {
1490
- "notes": "Expected values computed by an independent implementation written from the ONNX Runtime schema text alone, so this case checks the trusted reference as well as the kernels. SimplifiedLayerNormalization with no biases."
1491
  },
1492
  "attrs": { "K": 16, "N": 4, "bits": 4, "block_size": 8, "activation": "silu" },
1493
  "inputs": {
@@ -1538,7 +1539,7 @@
1538
  {
1539
  "name": "pinned_skipsum_gb_ub",
1540
  "provenance": {
1541
- "notes": "Expected values computed by an independent implementation written from the ONNX Runtime schema text alone, so this case checks the trusted reference as well as the kernels. SkipSimplifiedLayerNormalization with both biases and the residual-sum output."
1542
  },
1543
  "attrs": { "K": 16, "N": 4, "bits": 4, "block_size": 8, "activation": "silu" },
1544
  "inputs": {
@@ -1662,7 +1663,7 @@
1662
  {
1663
  "name": "norm_rows_past_one_tile",
1664
  "provenance": {
1665
- "notes": "Ten activation rows against a ROW_TILE of eight, so the row axis dispatches two groups and the second holds two real rows and six that clamp onto the last one. It is the only case where the store guard has anything to drop; every other multi-row case fits one group, where the guard cannot fire."
1666
  },
1667
  "attrs": { "K": 32, "N": 8, "bits": 4, "block_size": 16, "activation": "silu" },
1668
  "inputs": {
@@ -1698,6 +1699,348 @@
1698
  }
1699
  },
1700
  "outputs": { "yT": { "dtype": "float32", "shape": [10, 8], "tolerance": 0.0001, "relTolerance": 0.0001 } }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1701
  }
1702
  ]
1703
  }
 
1
  {
 
2
  "fixtureArrays": {
3
  "pinned_plain_gb_ub_input_aT": [0.9181, 1.1011, 1.115, 0.9522, 0.6333, 0.2043, -0.2696, -0.7143, -1.0586, -1.2453, -1.2411, -1.0421, -0.6753, -0.1947, 0.3263, 0.8068, 1.1707, 1.3587, 1.3385, 1.1098, 0.7051, 0.1852, -0.3707, -0.8762, -1.2521, -1.4384, -1.4042, -1.1526, -0.7208, -0.1744, 0.4029, 0.9218],
4
  "pinned_plain_gb_ub_input_gateBT": [52, 93, 210, 163, 160, 89, 30, 255, 204, 21, 42, 27, 184, 145, 246, 247, 100, 205, 130, 147, 208, 201, 206, 239, 252, 133, 218, 11, 232, 1, 166, 231],
 
1229
  },
1230
  {
1231
  "name": "block32_decode",
1232
+ "provenance": {
1233
+ "notes": "A single-row decode with K=64 and block_size=32 exercises two independently scaled 32-element blocks per output column."
1234
+ },
1235
  "attrs": { "K": 64, "N": 8, "bits": 4, "block_size": 32, "activation": "silu" },
1236
  "inputs": {
1237
  "aT": {
 
1432
  {
1433
  "name": "pinned_plain_gb_ub",
1434
  "provenance": {
1435
+ "notes": "Expected values are independently derived from the ONNX Runtime schema. With normalization omitted, both projections read A directly."
1436
  },
1437
  "attrs": { "K": 16, "N": 4, "bits": 4, "block_size": 8, "activation": "silu" },
1438
  "inputs": {
 
1488
  {
1489
  "name": "pinned_norm_nogb_noub",
1490
  "provenance": {
1491
+ "notes": "Expected values are independently derived from the ONNX Runtime schema. SimplifiedLayerNormalization runs without projection biases."
1492
  },
1493
  "attrs": { "K": 16, "N": 4, "bits": 4, "block_size": 8, "activation": "silu" },
1494
  "inputs": {
 
1539
  {
1540
  "name": "pinned_skipsum_gb_ub",
1541
  "provenance": {
1542
+ "notes": "Expected values are independently derived from the ONNX Runtime schema. SkipSimplifiedLayerNormalization uses both projection biases and returns the residual sum."
1543
  },
1544
  "attrs": { "K": 16, "N": 4, "bits": 4, "block_size": 8, "activation": "silu" },
1545
  "inputs": {
 
1663
  {
1664
  "name": "norm_rows_past_one_tile",
1665
  "provenance": {
1666
+ "notes": "Ten activation rows over an eight-row tile dispatch two groups. The second group has two live rows and six guarded rows clamped to the last valid input."
1667
  },
1668
  "attrs": { "K": 32, "N": 8, "bits": 4, "block_size": 16, "activation": "silu" },
1669
  "inputs": {
 
1699
  }
1700
  },
1701
  "outputs": { "yT": { "dtype": "float32", "shape": [10, 8], "tolerance": 0.0001, "relTolerance": 0.0001 } }
1702
+ },
1703
+ {
1704
+ "name": "decode_multi_trip_tails",
1705
+ "attrs": { "K": 300, "N": 5, "bits": 4, "block_size": 32, "activation": "silu" },
1706
+ "inputs": {
1707
+ "aT": {
1708
+ "dtype": "float32",
1709
+ "shape": [1, 300],
1710
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1711
+ },
1712
+ "normScaleT": {
1713
+ "dtype": "float32",
1714
+ "shape": [300],
1715
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
1716
+ },
1717
+ "gateBT": {
1718
+ "dtype": "uint8",
1719
+ "shape": [5, 10, 16],
1720
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
1721
+ },
1722
+ "gateScalesT": {
1723
+ "dtype": "float32",
1724
+ "shape": [5, 10],
1725
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
1726
+ },
1727
+ "gateBiasT": {
1728
+ "dtype": "float32",
1729
+ "shape": [5],
1730
+ "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.53, "scale": 0.3 }
1731
+ },
1732
+ "upBT": {
1733
+ "dtype": "uint8",
1734
+ "shape": [5, 10, 16],
1735
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
1736
+ },
1737
+ "upScalesT": {
1738
+ "dtype": "float32",
1739
+ "shape": [5, 10],
1740
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
1741
+ },
1742
+ "upBiasT": {
1743
+ "dtype": "float32",
1744
+ "shape": [5],
1745
+ "data": { "kind": "fillFloat32", "sinStep": 0.61, "cosStep": 0.17, "scale": 0.25 }
1746
+ }
1747
+ },
1748
+ "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001, "relTolerance": 0.0001 } },
1749
+ "provenance": {
1750
+ "notes": "Single row at K past two 128-code trips of the reduction-partitioned walk, with a partial final block and N not aligned to the decode column group."
1751
+ }
1752
+ },
1753
+ {
1754
+ "name": "decode_odd_k_skipsum",
1755
+ "attrs": { "K": 21, "N": 5, "bits": 4, "block_size": 32, "activation": "silu", "epsilon": 0.001 },
1756
+ "inputs": {
1757
+ "aT": {
1758
+ "dtype": "float32",
1759
+ "shape": [1, 21],
1760
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1761
+ },
1762
+ "skipT": {
1763
+ "dtype": "float32",
1764
+ "shape": [1, 21],
1765
+ "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
1766
+ },
1767
+ "normScaleT": {
1768
+ "dtype": "float32",
1769
+ "shape": [21],
1770
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
1771
+ },
1772
+ "gateBT": {
1773
+ "dtype": "uint8",
1774
+ "shape": [5, 1, 16],
1775
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
1776
+ },
1777
+ "gateScalesT": {
1778
+ "dtype": "float32",
1779
+ "shape": [5, 1],
1780
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
1781
+ },
1782
+ "upBT": {
1783
+ "dtype": "uint8",
1784
+ "shape": [5, 1, 16],
1785
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
1786
+ },
1787
+ "upScalesT": {
1788
+ "dtype": "float32",
1789
+ "shape": [5, 1],
1790
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
1791
+ }
1792
+ },
1793
+ "outputs": {
1794
+ "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001, "relTolerance": 0.0001 },
1795
+ "residualT": { "dtype": "float32", "shape": [1, 21], "tolerance": 0.000001, "relTolerance": 0.000001 }
1796
+ },
1797
+ "provenance": {
1798
+ "notes": "Single row with an odd K, so the last reduction pair has one live code; skip input and residual output."
1799
+ }
1800
+ },
1801
+ {
1802
+ "name": "bits8_decode_two_trips",
1803
+ "attrs": { "K": 160, "N": 8, "bits": 8, "block_size": 32, "activation": "silu" },
1804
+ "inputs": {
1805
+ "aT": {
1806
+ "dtype": "float32",
1807
+ "shape": [1, 160],
1808
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1809
+ },
1810
+ "normScaleT": {
1811
+ "dtype": "float32",
1812
+ "shape": [160],
1813
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
1814
+ },
1815
+ "gateBT": {
1816
+ "dtype": "uint8",
1817
+ "shape": [8, 5, 32],
1818
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
1819
+ },
1820
+ "gateScalesT": {
1821
+ "dtype": "float32",
1822
+ "shape": [8, 5],
1823
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
1824
+ },
1825
+ "upBT": {
1826
+ "dtype": "uint8",
1827
+ "shape": [8, 5, 32],
1828
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
1829
+ },
1830
+ "upScalesT": {
1831
+ "dtype": "float32",
1832
+ "shape": [8, 5],
1833
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
1834
+ }
1835
+ },
1836
+ "outputs": { "yT": { "dtype": "float32", "shape": [1, 8], "tolerance": 0.0001, "relTolerance": 0.0001 } },
1837
+ "provenance": { "notes": "Single row at 8 bits, where a code pair spans two stored slots, over two trips." }
1838
+ },
1839
+ {
1840
+ "name": "plain_decode_two_trips",
1841
+ "attrs": { "K": 200, "N": 6, "bits": 4, "block_size": 32, "activation": "silu" },
1842
+ "inputs": {
1843
+ "aT": {
1844
+ "dtype": "float32",
1845
+ "shape": [1, 200],
1846
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1847
+ },
1848
+ "gateBT": {
1849
+ "dtype": "uint8",
1850
+ "shape": [6, 7, 16],
1851
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
1852
+ },
1853
+ "gateScalesT": {
1854
+ "dtype": "float32",
1855
+ "shape": [6, 7],
1856
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
1857
+ },
1858
+ "gateBiasT": {
1859
+ "dtype": "float32",
1860
+ "shape": [6],
1861
+ "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.53, "scale": 0.3 }
1862
+ },
1863
+ "upBT": {
1864
+ "dtype": "uint8",
1865
+ "shape": [6, 7, 16],
1866
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
1867
+ },
1868
+ "upScalesT": {
1869
+ "dtype": "float32",
1870
+ "shape": [6, 7],
1871
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
1872
+ },
1873
+ "upBiasT": {
1874
+ "dtype": "float32",
1875
+ "shape": [6],
1876
+ "data": { "kind": "fillFloat32", "sinStep": 0.61, "cosStep": 0.17, "scale": 0.25 }
1877
+ }
1878
+ },
1879
+ "outputs": { "yT": { "dtype": "float32", "shape": [1, 6], "tolerance": 0.0001, "relTolerance": 0.0001 } },
1880
+ "provenance": {
1881
+ "notes": "Single row without normalization, so the reduction-partitioned walk reads the raw activation."
1882
+ }
1883
+ },
1884
+ {
1885
+ "name": "block32_decode_eight_columns_n13",
1886
+ "provenance": {
1887
+ "notes": "K=64 at block 32 is two sixteen-byte vector groups per column, at most one per lane, so the decode walk covers eight columns per workgroup; N=13 leaves the second workgroup with a full first column group and a one-column second group, exercising the guards on columns 4..7."
1888
+ },
1889
+ "attrs": { "K": 64, "N": 13, "bits": 4, "block_size": 32, "activation": "silu" },
1890
+ "inputs": {
1891
+ "aT": {
1892
+ "dtype": "float32",
1893
+ "shape": [1, 64],
1894
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1895
+ },
1896
+ "skipT": {
1897
+ "dtype": "float32",
1898
+ "shape": [1, 64],
1899
+ "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
1900
+ },
1901
+ "normScaleT": {
1902
+ "dtype": "float32",
1903
+ "shape": [64],
1904
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
1905
+ },
1906
+ "gateBT": {
1907
+ "dtype": "uint8",
1908
+ "shape": [13, 2, 16],
1909
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
1910
+ },
1911
+ "gateScalesT": {
1912
+ "dtype": "float32",
1913
+ "shape": [13, 2],
1914
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
1915
+ },
1916
+ "gateBiasT": {
1917
+ "dtype": "float32",
1918
+ "shape": [13],
1919
+ "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.53, "scale": 0.3 }
1920
+ },
1921
+ "upBT": {
1922
+ "dtype": "uint8",
1923
+ "shape": [13, 2, 16],
1924
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
1925
+ },
1926
+ "upScalesT": {
1927
+ "dtype": "float32",
1928
+ "shape": [13, 2],
1929
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
1930
+ },
1931
+ "upBiasT": {
1932
+ "dtype": "float32",
1933
+ "shape": [13],
1934
+ "data": { "kind": "fillFloat32", "sinStep": 0.61, "cosStep": 0.17, "scale": 0.25 }
1935
+ }
1936
+ },
1937
+ "outputs": {
1938
+ "yT": { "dtype": "float32", "shape": [1, 13], "tolerance": 0.0001, "relTolerance": 0.0001 },
1939
+ "residualT": { "dtype": "float32", "shape": [1, 64], "tolerance": 0.000001, "relTolerance": 0.000001 }
1940
+ }
1941
+ },
1942
+ {
1943
+ "name": "block32_decode_eight_columns_n5_plain",
1944
+ "provenance": {
1945
+ "notes": "Eight-column decode walk without the norm pass (plain activations, no skip): N=5 is one workgroup whose second column group holds a single live column."
1946
+ },
1947
+ "attrs": { "K": 64, "N": 5, "bits": 4, "block_size": 32, "activation": "silu" },
1948
+ "inputs": {
1949
+ "aT": {
1950
+ "dtype": "float32",
1951
+ "shape": [1, 64],
1952
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1953
+ },
1954
+ "gateBT": {
1955
+ "dtype": "uint8",
1956
+ "shape": [5, 2, 16],
1957
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
1958
+ },
1959
+ "gateScalesT": {
1960
+ "dtype": "float32",
1961
+ "shape": [5, 2],
1962
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
1963
+ },
1964
+ "gateBiasT": {
1965
+ "dtype": "float32",
1966
+ "shape": [5],
1967
+ "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.53, "scale": 0.3 }
1968
+ },
1969
+ "upBT": {
1970
+ "dtype": "uint8",
1971
+ "shape": [5, 2, 16],
1972
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
1973
+ },
1974
+ "upScalesT": {
1975
+ "dtype": "float32",
1976
+ "shape": [5, 2],
1977
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
1978
+ },
1979
+ "upBiasT": {
1980
+ "dtype": "float32",
1981
+ "shape": [5],
1982
+ "data": { "kind": "fillFloat32", "sinStep": 0.61, "cosStep": 0.17, "scale": 0.25 }
1983
+ }
1984
+ },
1985
+ "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001, "relTolerance": 0.0001 } }
1986
+ },
1987
+ {
1988
+ "name": "block32_decode_k2080_two_groups_per_lane",
1989
+ "provenance": {
1990
+ "notes": "K=2080 at block 32 is 65 sixteen-byte vector groups per column, more than the 64-lane decode workgroup, so the vector walk keeps both projections in every lane (no lane split) and lane 0 takes a second group; the staged norm covers 2080 activations."
1991
+ },
1992
+ "attrs": { "K": 2080, "N": 8, "bits": 4, "block_size": 32, "activation": "silu" },
1993
+ "inputs": {
1994
+ "aT": {
1995
+ "dtype": "float32",
1996
+ "shape": [1, 2080],
1997
+ "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
1998
+ },
1999
+ "skipT": {
2000
+ "dtype": "float32",
2001
+ "shape": [1, 2080],
2002
+ "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
2003
+ },
2004
+ "normScaleT": {
2005
+ "dtype": "float32",
2006
+ "shape": [2080],
2007
+ "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
2008
+ },
2009
+ "gateBT": {
2010
+ "dtype": "uint8",
2011
+ "shape": [8, 65, 16],
2012
+ "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
2013
+ },
2014
+ "gateScalesT": {
2015
+ "dtype": "float32",
2016
+ "shape": [8, 65],
2017
+ "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
2018
+ },
2019
+ "gateBiasT": {
2020
+ "dtype": "float32",
2021
+ "shape": [8],
2022
+ "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.53, "scale": 0.3 }
2023
+ },
2024
+ "upBT": {
2025
+ "dtype": "uint8",
2026
+ "shape": [8, 65, 16],
2027
+ "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
2028
+ },
2029
+ "upScalesT": {
2030
+ "dtype": "float32",
2031
+ "shape": [8, 65],
2032
+ "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
2033
+ },
2034
+ "upBiasT": {
2035
+ "dtype": "float32",
2036
+ "shape": [8],
2037
+ "data": { "kind": "fillFloat32", "sinStep": 0.61, "cosStep": 0.17, "scale": 0.25 }
2038
+ }
2039
+ },
2040
+ "outputs": {
2041
+ "yT": { "dtype": "float32", "shape": [1, 8], "tolerance": 0.001, "relTolerance": 0.001 },
2042
+ "residualT": { "dtype": "float32", "shape": [1, 2080], "tolerance": 0.000001, "relTolerance": 0.000001 }
2043
+ }
2044
  }
2045
  ]
2046
  }