Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.GemmFastGelu

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
26fd85e
·
verified ·
1 Parent(s): fe90e05

sync c928d21e6cc1

Browse files
Files changed (3) hide show
  1. build/webgpu/gemm-fast-gelu.wgsl.jinja +5 -1
  2. build/webgpu/gemm-subgroup-matrix.wgsl.jinja +13 -7
  3. build/webgpu/metadata.json +4 -4
build/webgpu/gemm-fast-gelu.wgsl.jinja CHANGED
@@ -35,7 +35,11 @@ fn fused_act(v: f32) -> f32 {
35
  return 1.0 / (1.0 + exp(-v));
36
  {% endif %}
37
  {% if actMode == "Tanh" %}
38
- return tanh(v);
 
 
 
 
39
  {% endif %}
40
  {% if actMode == "HardSigmoid" %}
41
  return clamp(f32({{ actA }}) * v + f32({{ actB }}), 0.0, 1.0);
 
35
  return 1.0 / (1.0 + exp(-v));
36
  {% endif %}
37
  {% if actMode == "Tanh" %}
38
+ // Clamp first: the hardware tanh evaluates (e^2v - 1)/(e^2v + 1), which is
39
+ // Inf/Inf = NaN once e^2v overflows (v > ~44). A GEMM accumulator reaches that
40
+ // easily because it grows with K. tanh is already +/-1 to full f32 precision by
41
+ // |v| ~ 9, so the clamp changes no representable result.
42
+ return tanh(clamp(v, -10.0, 10.0));
43
  {% endif %}
44
  {% if actMode == "HardSigmoid" %}
45
  return clamp(f32({{ actA }}) * v + f32({{ actB }}), 0.0, 1.0);
build/webgpu/gemm-subgroup-matrix.wgsl.jinja CHANGED
@@ -12,6 +12,8 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
12
 
13
  {% set operandScalar = mmaScalar %}
14
  {% set accScalar = "f32" %}
 
 
15
  {% set tileMValue = tileM if tileM is defined else 32 %}
16
  {% set tileNValue = tileN if tileN is defined else 64 %}
17
  {% set tileKValue = tileK if tileK is defined else 32 %}
@@ -55,7 +57,11 @@ fn fused_act(v: f32) -> f32 {
55
  return 1.0 / (1.0 + exp(-v));
56
  {% endif %}
57
  {% if actMode == "Tanh" %}
58
- return tanh(v);
 
 
 
 
59
  {% endif %}
60
  {% if actMode == "HardSigmoid" %}
61
  return clamp(f32({{ actA }}) * v + f32({{ actB }}), 0.0, 1.0);
@@ -139,8 +145,8 @@ fn loadSHMB(tile_base: u32, k_idx: u32, row: u32, c_idx: u32, bank_offset: u32)
139
  fn storeOutput(offset: u32{% if needsColBase %}, col_base: u32{% endif %}, row: u32, col: u32, src_slot: u32, row_limit: i32) {
140
  if (row_limit > 0 && row < u32(row_limit)) {
141
  let col2: u32 = col + 1u;
142
- y[offset + row * OUT_F + col] = {{ yScalar }}({{ epiOpen }}scratch[src_slot][row * 8u + col]{{ BIAS_1 }}{{ epiClose }});
143
- y[offset + row * OUT_F + col2] = {{ yScalar }}({{ epiOpen }}scratch[src_slot][row * 8u + col2]{{ BIAS_2 }}{{ epiClose }});
144
  }
145
  }
146
 
@@ -182,7 +188,7 @@ fn main(
182
  {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
183
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
184
  {% for m in range(dynamicATiles) %}
185
- var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, false, TILE_K);
186
  {% endfor %}
187
 
188
  {% if bTransposed %}
@@ -192,9 +198,9 @@ fn main(
192
  {% endif %}
193
  {% for n in range(bTilesPerSubgroup) %}
194
  {% if bTransposed %}
195
- var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + {{ n * 8 }}u, false, TILE_COLS);
196
  {% else %}
197
- var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, true, TILE_K);
198
  {% endif %}
199
  {% endfor %}
200
 
@@ -218,7 +224,7 @@ fn main(
218
  let bank: u32 = subtile_id * {{ aTilesPerSubgroup * bTilesPerSubgroup }}u;
219
  {% for m in range(aTilesPerSubgroup) %}
220
  {% for n in range(bTilesPerSubgroup) %}
221
- subgroupMatrixStore(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matC{{ m }}{{ n }}, false, 8u);
222
  {% endfor %}
223
  {% endfor %}
224
  workgroupBarrier();
 
12
 
13
  {% set operandScalar = mmaScalar %}
14
  {% set accScalar = "f32" %}
15
+ {% set OUT = outputBuffer if outputBuffer is defined else "y" %}
16
+ {% set OUT_SCALAR = outScalar if outScalar is defined else yScalar %}
17
  {% set tileMValue = tileM if tileM is defined else 32 %}
18
  {% set tileNValue = tileN if tileN is defined else 64 %}
19
  {% set tileKValue = tileK if tileK is defined else 32 %}
 
57
  return 1.0 / (1.0 + exp(-v));
58
  {% endif %}
59
  {% if actMode == "Tanh" %}
60
+ // Clamp first: the hardware tanh evaluates (e^2v - 1)/(e^2v + 1), which is
61
+ // Inf/Inf = NaN once e^2v overflows (v > ~44). A GEMM accumulator reaches that
62
+ // easily because it grows with K. tanh is already +/-1 to full f32 precision by
63
+ // |v| ~ 9, so the clamp changes no representable result.
64
+ return tanh(clamp(v, -10.0, 10.0));
65
  {% endif %}
66
  {% if actMode == "HardSigmoid" %}
67
  return clamp(f32({{ actA }}) * v + f32({{ actB }}), 0.0, 1.0);
 
145
  fn storeOutput(offset: u32{% if needsColBase %}, col_base: u32{% endif %}, row: u32, col: u32, src_slot: u32, row_limit: i32) {
146
  if (row_limit > 0 && row < u32(row_limit)) {
147
  let col2: u32 = col + 1u;
148
+ {{ OUT }}[offset + row * OUT_F + col] = {{ OUT_SCALAR }}({{ epiOpen }}scratch[src_slot][row * 8u + col]{{ BIAS_1 }}{{ epiClose }});
149
+ {{ OUT }}[offset + row * OUT_F + col2] = {{ OUT_SCALAR }}({{ epiOpen }}scratch[src_slot][row * 8u + col2]{{ BIAS_2 }}{{ epiClose }});
150
  }
151
  }
152
 
 
188
  {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
189
  let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
190
  {% for m in range(dynamicATiles) %}
191
+ var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
192
  {% endfor %}
193
 
194
  {% if bTransposed %}
 
198
  {% endif %}
199
  {% for n in range(bTilesPerSubgroup) %}
200
  {% if bTransposed %}
201
+ var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
202
  {% else %}
203
+ var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, TILE_K);
204
  {% endif %}
205
  {% endfor %}
206
 
 
224
  let bank: u32 = subtile_id * {{ aTilesPerSubgroup * bTilesPerSubgroup }}u;
225
  {% for m in range(aTilesPerSubgroup) %}
226
  {% for n in range(bTilesPerSubgroup) %}
227
+ subgroupMatrixStore<row_major>(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matC{{ m }}{{ n }}, 8u);
228
  {% endfor %}
229
  {% endfor %}
230
  workgroupBarrier();
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "com.microsoft.GemmFastGelu",
3
- "id": "_com_microsoft_gemmfastgelu_webgpu_372327f",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -8,12 +8,12 @@
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "CPFUVoaXFHE1nKK5oZMKmB0LohpzyltnfuRq/70nJAk=",
11
- "gemm-fast-gelu.wgsl.jinja": "IkB53JqVNn9Xezw0YcKYDvm9IwDvoZNnSeA/UY4Eci4=",
12
- "gemm-subgroup-matrix.wgsl.jinja": "2PtV5JZcKiR+CKWoCbReeevzatqfe9l/O6U5wwbDgoU=",
13
  "manifest.json": "bcpv0SoGGVwU23+3nF7YJ2BSX80r64AInCEZgAuCRoI=",
14
  "test.json": "WNyAFlYnnkTg3N4r9ja8Kh2HwlO8jr7EncOBBa9oI48="
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
18
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GemmFastGelu" }
19
  }
 
1
  {
2
  "name": "com.microsoft.GemmFastGelu",
3
+ "id": "_com_microsoft_gemmfastgelu_webgpu_720afad",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
8
  "algorithm": "sha256",
9
  "files": {
10
  "bench.json": "CPFUVoaXFHE1nKK5oZMKmB0LohpzyltnfuRq/70nJAk=",
11
+ "gemm-fast-gelu.wgsl.jinja": "XL5rMA9oUlUJZCMHq2PZAzmORxFZoz/xa5Mo05/aK3k=",
12
+ "gemm-subgroup-matrix.wgsl.jinja": "JJZezmRCzycCZ5Y6RTmbisasCFLbVAgwkSkV4Sct7eY=",
13
  "manifest.json": "bcpv0SoGGVwU23+3nF7YJ2BSX80r64AInCEZgAuCRoI=",
14
  "test.json": "WNyAFlYnnkTg3N4r9ja8Kh2HwlO8jr7EncOBBa9oI48="
15
  }
16
  },
17
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
18
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GemmFastGelu" }
19
  }