Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.BiasSplitGelu

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
a44ec85
·
verified ·
1 Parent(s): d1568b9

sync 6fdf6301e2bb

Browse files
Files changed (5) hide show
  1. README.md +2 -2
  2. build/webgpu/bias-split-gelu.wgsl.jinja +5 -31
  3. build/webgpu/manifest.json +13 -20
  4. build/webgpu/metadata.json +6 -6
  5. build/webgpu/test.json +3 -3
README.md CHANGED
@@ -40,13 +40,13 @@ See the [ONNX Runtime `BiasSplitGelu` contrib-operator spec](https://github.com/
40
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
41
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
42
  - [`test.json`](build/webgpu/test.json) — correctness cases
43
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
44
  - [`bias-split-gelu.wgsl.jinja`](build/webgpu/bias-split-gelu.wgsl.jinja)
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
  ```sh
49
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
50
  ```
51
 
52
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
40
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
41
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
42
  - [`test.json`](build/webgpu/test.json) — correctness cases
43
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
44
  - [`bias-split-gelu.wgsl.jinja`](build/webgpu/bias-split-gelu.wgsl.jinja)
45
 
46
  ## Use with `@huggingface/kernels`
47
 
48
  ```sh
49
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
50
  ```
51
 
52
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/bias-split-gelu.wgsl.jinja CHANGED
@@ -1,36 +1,11 @@
1
- {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
2
- {% if note == "dispatch-limit" %}
3
- // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
- // per-axis workgroup fold width (outputs > 16.7M elements).
5
- {% elif note == "limit" %}
6
- // 2D-folded flat index: gid.y carries the high bits past the dispatch's
7
- // per-axis workgroup fold width.
8
- {% elif note == "device-axis" %}
9
- // The flat dispatch is folded across x/y at a fixed per-axis workgroup
10
- // width; gid.y carries the high portion of the output index.
11
- {% elif note == "vec4-limit" %}
12
- // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
13
- // per-axis workgroup fold width (the dispatch caps x and spills into y).
14
- {% elif note == "element-limit" %}
15
- // 2D-folded flat element index: gid.y carries the high bits past the
16
- // dispatch's per-axis workgroup fold width.
17
- {% elif note == "dispatch" %}
18
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
19
  // per-axis workgroup fold width.
20
- {% endif %}
21
- {% if bound == "" %}
22
- let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
23
- {%- elif guardInline %}
24
- let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
25
- if ({{ name }} >= {{ bound }}) { return; }
26
- {%- else %}
27
- let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
28
  if ({{ name }} >= {{ bound }}) {
29
  return;
30
- }
31
- {%- endif %}
32
- {% endmacro %}
33
-
34
  {{ env.wgsl.resourceDeclarations }}
35
 
36
  // com.microsoft.BiasSplitGelu GEGLU gate:
@@ -54,7 +29,6 @@ fn erf_approx(x: f32) -> f32 {
54
  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-(ax * ax));
55
  return sign * y;
56
  }
57
-
58
  fn gelu_value(v: f32) -> f32 {
59
  return 0.5 * v * (1.0 + erf_approx(v * 0.7071067811865476));
60
  }
@@ -67,7 +41,7 @@ const GROUPS_PER_ROW: u32 = (HIDDEN + 3u) / 4u;
67
  {% endif %}
68
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
69
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
70
- {{ flat_index_2d() }}
71
  {% if vec4Tail %}
72
  let row = i / GROUPS_PER_ROW;
73
  let col = (i % GROUPS_PER_ROW) * 4u;
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
  // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
 
 
 
 
 
 
 
6
  if ({{ name }} >= {{ bound }}) {
7
  return;
8
+ }{% endmacro %}
 
 
 
9
  {{ env.wgsl.resourceDeclarations }}
10
 
11
  // com.microsoft.BiasSplitGelu GEGLU gate:
 
29
  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-(ax * ax));
30
  return sign * y;
31
  }
 
32
  fn gelu_value(v: f32) -> f32 {
33
  return 0.5 * v * (1.0 + erf_approx(v * 0.7071067811865476));
34
  }
 
41
  {% endif %}
42
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
43
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
44
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE) }}
45
  {% if vec4Tail %}
46
  let row = i / GROUPS_PER_ROW;
47
  let col = (i % GROUPS_PER_ROW) * 4u;
build/webgpu/manifest.json CHANGED
@@ -8,12 +8,11 @@
8
  "tunables": { "WORKGROUP_SIZE": { "default": 64 } },
9
  "when": ["ranks.X == 3", "ranks.Y == 3", "dim(shapes.Y, 0) == dim(shapes.X, 0)", "dim(shapes.Y, 1) == dim(shapes.X, 1)", "dim(shapes.X, 2) % 2 == 0", "dim(shapes.Y, 2) == dim(shapes.X, 2) / 2", "f16Ok(dtypes.T)", "ranks.bias == 1", "dim(shapes.bias, 0) == dim(shapes.X, 2)"],
10
  "bindings": {
11
- "x": { "arg": "X", "buffer": "read-only-storage", "elementType": "$scalar" },
12
- "bias": { "buffer": "read-only-storage", "elementType": "$scalar", "length": "$twice" },
13
- "y_2": { "arg": "Y", "name": "y", "buffer": "storage", "elementType": "$scalar" },
14
- "params_2": {
15
  "name": "params",
16
- "buffer": "uniform",
17
  "struct": [
18
  {
19
  "name": "count",
@@ -22,12 +21,8 @@
22
  }
23
  ]
24
  },
25
- "bias_2": { "name": "bias", "buffer": "read-only-storage", "elementType": "$scalar" },
26
- "params_3": {
27
- "name": "params",
28
- "buffer": "uniform",
29
- "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }]
30
- }
31
  },
32
  "variants": [
33
  {
@@ -35,7 +30,6 @@
35
  "priority": 30,
36
  "when": ["numel(shapes.Y) > 0", "numel(shapes.Y) % 4 == 0", "dim(shapes.Y, 2) % 4 == 0"],
37
  "derive": {
38
- "scalar": "dtypes.T",
39
  "vec4": true,
40
  "vec4Tail": false,
41
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
@@ -45,7 +39,7 @@
45
  "passes": [
46
  {
47
  "id": "main",
48
- "name": "BiasSplitGelu.vec4",
49
  "shader": "bias-split-gelu.wgsl.jinja",
50
  "bindings": [
51
  "x",
@@ -66,7 +60,6 @@
66
  "priority": 20,
67
  "when": ["numel(shapes.Y) > 0"],
68
  "derive": {
69
- "scalar": "dtypes.T",
70
  "vec4": false,
71
  "vec4Tail": true,
72
  "hidden": "dim(shapes.X, ranks.X - 1) / 2 if dim(shapes.X, ranks.X - 1) > 0 else 1",
@@ -75,9 +68,9 @@
75
  "passes": [
76
  {
77
  "id": "main",
78
- "name": "BiasSplitGelu.vec4Tail",
79
  "shader": "bias-split-gelu.wgsl.jinja",
80
- "bindings": ["x", "bias", "y_2", "params_2"],
81
  "dispatch": {
82
  "x": "min(ceilDiv(((numel(shapes.Y) / (dim(shapes.X, ranks.X - 1) / 2)) * ceilDiv(dim(shapes.X, ranks.X - 1) / 2, 4)), (tunables.WORKGROUP_SIZE)), 65535)",
83
  "y": "ceilDiv(ceilDiv(((numel(shapes.Y) / (dim(shapes.X, ranks.X - 1) / 2)) * ceilDiv(dim(shapes.X, ranks.X - 1) / 2, 4)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -91,7 +84,6 @@
91
  "priority": 0,
92
  "when": ["true"],
93
  "derive": {
94
- "scalar": "dtypes.T",
95
  "vec4": false,
96
  "vec4Tail": false,
97
  "hidden": "dim(shapes.X, ranks.X - 1) / 2 if dim(shapes.X, ranks.X - 1) > 0 else 1",
@@ -100,9 +92,9 @@
100
  "passes": [
101
  {
102
  "id": "main",
103
- "name": "BiasSplitGelu.scalar",
104
  "shader": "bias-split-gelu.wgsl.jinja",
105
- "bindings": ["x", "bias_2", "y_2", "params_3"],
106
  "dispatch": {
107
  "x": "min(ceilDiv((numel(shapes.Y)), (tunables.WORKGROUP_SIZE)), 65535)",
108
  "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -111,5 +103,6 @@
111
  }
112
  ]
113
  }
114
- ]
 
115
  }
 
8
  "tunables": { "WORKGROUP_SIZE": { "default": 64 } },
9
  "when": ["ranks.X == 3", "ranks.Y == 3", "dim(shapes.Y, 0) == dim(shapes.X, 0)", "dim(shapes.Y, 1) == dim(shapes.X, 1)", "dim(shapes.X, 2) % 2 == 0", "dim(shapes.Y, 2) == dim(shapes.X, 2) / 2", "f16Ok(dtypes.T)", "ranks.bias == 1", "dim(shapes.bias, 0) == dim(shapes.X, 2)"],
10
  "bindings": {
11
+ "x": { "arg": "X", "elementType": "$scalar" },
12
+ "bias": { "elementType": "$scalar", "length": "$twice" },
13
+ "y_y": { "arg": "Y", "name": "y", "elementType": "$scalar" },
14
+ "params_main": {
15
  "name": "params",
 
16
  "struct": [
17
  {
18
  "name": "count",
 
21
  }
22
  ]
23
  },
24
+ "bias_main": { "name": "bias", "elementType": "$scalar" },
25
+ "params_count": { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.Y)" }] }
 
 
 
 
26
  },
27
  "variants": [
28
  {
 
30
  "priority": 30,
31
  "when": ["numel(shapes.Y) > 0", "numel(shapes.Y) % 4 == 0", "dim(shapes.Y, 2) % 4 == 0"],
32
  "derive": {
 
33
  "vec4": true,
34
  "vec4Tail": false,
35
  "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
 
39
  "passes": [
40
  {
41
  "id": "main",
42
+ "name": "BiasSplitGelu.Vec4",
43
  "shader": "bias-split-gelu.wgsl.jinja",
44
  "bindings": [
45
  "x",
 
60
  "priority": 20,
61
  "when": ["numel(shapes.Y) > 0"],
62
  "derive": {
 
63
  "vec4": false,
64
  "vec4Tail": true,
65
  "hidden": "dim(shapes.X, ranks.X - 1) / 2 if dim(shapes.X, ranks.X - 1) > 0 else 1",
 
68
  "passes": [
69
  {
70
  "id": "main",
71
+ "name": "BiasSplitGelu.Vec4Tail",
72
  "shader": "bias-split-gelu.wgsl.jinja",
73
+ "bindings": ["x", "bias", "y_y", "params_main"],
74
  "dispatch": {
75
  "x": "min(ceilDiv(((numel(shapes.Y) / (dim(shapes.X, ranks.X - 1) / 2)) * ceilDiv(dim(shapes.X, ranks.X - 1) / 2, 4)), (tunables.WORKGROUP_SIZE)), 65535)",
76
  "y": "ceilDiv(ceilDiv(((numel(shapes.Y) / (dim(shapes.X, ranks.X - 1) / 2)) * ceilDiv(dim(shapes.X, ranks.X - 1) / 2, 4)), (tunables.WORKGROUP_SIZE)), 65535)",
 
84
  "priority": 0,
85
  "when": ["true"],
86
  "derive": {
 
87
  "vec4": false,
88
  "vec4Tail": false,
89
  "hidden": "dim(shapes.X, ranks.X - 1) / 2 if dim(shapes.X, ranks.X - 1) > 0 else 1",
 
92
  "passes": [
93
  {
94
  "id": "main",
95
+ "name": "BiasSplitGelu.Scalar",
96
  "shader": "bias-split-gelu.wgsl.jinja",
97
+ "bindings": ["x", "bias_main", "y_y", "params_count"],
98
  "dispatch": {
99
  "x": "min(ceilDiv((numel(shapes.Y)), (tunables.WORKGROUP_SIZE)), 65535)",
100
  "y": "ceilDiv(ceilDiv((numel(shapes.Y)), (tunables.WORKGROUP_SIZE)), 65535)",
 
103
  }
104
  ]
105
  }
106
+ ],
107
+ "derive": { "scalar": "dtypes.T" }
108
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "com.microsoft.BiasSplitGelu",
3
- "id": "_com_microsoft_biassplitgelu_webgpu_c7117c6",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,14 +8,14 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "9D5pXr6GQoiKWOaIZonasPZV750s7EBoeCc18Fjh6LE=",
11
- "bias-split-gelu.wgsl.jinja": "NdmUm+O7ZCH0nHgT9xuUVfsMcX0TKQLoz2edprOP4Eg=",
12
- "manifest.json": "RhauSBIjgU0p1ri0W74BHIABvKxAMnSg9MfBB++0c0E=",
13
- "test.json": "jeqt2QA6oOrulPiXyZ1nXe7lvFKPg1me9DSnpzLx7lk="
14
  }
15
  },
16
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
17
  "webgpu": {
18
- "manifestSpec": "2.0",
19
  "variants": {
20
  "vec4": ["bias-split-gelu.wgsl.jinja"],
21
  "vec4_tail": ["bias-split-gelu.wgsl.jinja"],
 
1
  {
2
  "name": "com.microsoft.BiasSplitGelu",
3
+ "id": "_com_microsoft_biassplitgelu_webgpu_8c2c282",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "9D5pXr6GQoiKWOaIZonasPZV750s7EBoeCc18Fjh6LE=",
11
+ "bias-split-gelu.wgsl.jinja": "e+SSgSOxw31lOAadfmdFjmjkSgsQ3LMMTvWCzJ32SqA=",
12
+ "manifest.json": "aqN8b2kpzusQElQ1g5RAKFuEGVgCEVx9LkPmxHEOpeY=",
13
+ "test.json": "48AxMaJkqHxGNEXQV5nWesTxizSJr+GET4DxjeCcBmI="
14
  }
15
  },
16
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
17
  "webgpu": {
18
+ "manifestSpec": "2.1",
19
  "variants": {
20
  "vec4": ["bias-split-gelu.wgsl.jinja"],
21
  "vec4_tail": ["bias-split-gelu.wgsl.jinja"],
build/webgpu/test.json CHANGED
@@ -144,7 +144,7 @@
144
  "name": "split_gelu_d4_vec4_path",
145
  "provenance": {
146
  "source": "onnxruntime/test/contrib_ops/bias_split_gelu_op_test.cc",
147
- "notes": "D=4 (input hidden 8) so the output hidden axis is a multiple of 4 and the vec4 manifest variant is exercised. Expected values from the same erf-form GEGLU gate computed at f64 then f32-rounded."
148
  },
149
  "inputs": {
150
  "X": {
@@ -208,7 +208,7 @@
208
  {
209
  "name": "f16_vec4_shader_validation",
210
  "provenance": {
211
- "notes": "Non-empty f16 D=4 case forces the vec4 variant; the empty f16 case cannot catch shader construction errors because it returns before compiling meaningful lane loads."
212
  },
213
  "inputs": {
214
  "X": {
@@ -227,7 +227,7 @@
227
  {
228
  "name": "scalar_path_dispatch_fold_over_16M",
229
  "provenance": {
230
- "notes": "Output width 6 is not divisible by four, selecting the scalar variant. With 16,776,966 output elements, the two-dimensional dispatch fold must reconstruct the linear index from both grid axes and guard the final partial row."
231
  },
232
  "inputs": {
233
  "X": {
 
144
  "name": "split_gelu_d4_vec4_path",
145
  "provenance": {
146
  "source": "onnxruntime/test/contrib_ops/bias_split_gelu_op_test.cc",
147
+ "notes": "D=4 (input hidden 8) gives a four-aligned output hidden axis. Expected values from the same erf-form GEGLU gate computed at f64 then f32-rounded."
148
  },
149
  "inputs": {
150
  "X": {
 
208
  {
209
  "name": "f16_vec4_shader_validation",
210
  "provenance": {
211
+ "notes": "A non-empty float16 input (channel size 8, output size 4) checks BiasSplitGelu end-to-end; this file's empty_zero_dim_f16 case has a zero-sized dimension and returns before touching data, so it can't validate the computation alone."
212
  },
213
  "inputs": {
214
  "X": {
 
227
  {
228
  "name": "scalar_path_dispatch_fold_over_16M",
229
  "provenance": {
230
+ "notes": "An output width of six and 16,776,966 elements exercise unaligned channels and reconstruction of a folded two-dimensional dispatch grid."
231
  },
232
  "inputs": {
233
  "X": {