Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.QMoE

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
3626d99
·
verified ·
1 Parent(s): 793c2c2

sync c928d21e6cc1

Browse files
Files changed (4) hide show
  1. build/webgpu/manifest.json +5 -2
  2. build/webgpu/metadata.json +5 -5
  3. build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja +11 -8
  4. build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja +12 -9
build/webgpu/manifest.json CHANGED
@@ -104,6 +104,10 @@
104
  },
105
  "derive": {
106
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
 
 
 
 
107
  "hiddenSize": "dim(shapes.input, ranks.input - 1)",
108
  "numTokens": "numel(shapes.input) / max(1, hiddenSize)",
109
  "topK": "attrs.k",
@@ -574,7 +578,7 @@
574
  {
575
  "id": "quant_grouped_sgmat_routed",
576
  "priority": 32,
577
- "when": ["groupedSgmatOk"],
578
  "constants": {
579
  "groupTileM": "groupTileM",
580
  "groupRouteWorkgroup": "tunables.groupRouteWorkgroup",
@@ -642,7 +646,6 @@
642
  "requires": {
643
  "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
644
  "limits": { "maxComputeWorkgroupStorageSize": 16512 },
645
- "subgroupMinSize": 32,
646
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
647
  }
648
  },
 
104
  },
105
  "derive": {
106
  "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
107
+ "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
108
+ "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
109
+ "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
110
+ "wave32Effective": "wave32Adapter or pinSubgroupSize32",
111
  "hiddenSize": "dim(shapes.input, ranks.input - 1)",
112
  "numTokens": "numel(shapes.input) / max(1, hiddenSize)",
113
  "topK": "attrs.k",
 
578
  {
579
  "id": "quant_grouped_sgmat_routed",
580
  "priority": 32,
581
+ "when": ["groupedSgmatOk", "wave32Effective"],
582
  "constants": {
583
  "groupTileM": "groupTileM",
584
  "groupRouteWorkgroup": "tunables.groupRouteWorkgroup",
 
646
  "requires": {
647
  "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
648
  "limits": { "maxComputeWorkgroupStorageSize": 16512 },
 
649
  "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
650
  }
651
  },
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "com.microsoft.QMoE",
3
- "id": "_com_microsoft_qmoe_webgpu_9a12372",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -9,14 +9,14 @@
9
  "files": {
10
  "bench.json": "zBVupnayA0Ggc4y737P9tzqyM/q/N8F4j+FEJXz6Wks=",
11
  "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
12
- "manifest.json": "hZMDKX6e3562/hNqPc5QSpMZ7BRW0wcSrH63e8VTdM4=",
13
  "qmoe-fc1-activation-gemv.wgsl.jinja": "UvP8Vyx1E9+/+I09w/z9iB6h6MHrSrOmjKfIelyIS0M=",
14
- "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "bRKPszdbDTG0DlHeiiQ1lxyjJUJz0JVt9WZDsv8KVg8=",
15
  "qmoe-fc1-activation-grouped.wgsl.jinja": "y2vGeNeXE5YmNrM2WrXogyF/fTvQ/nqYR4uWBTDvpPY=",
16
  "qmoe-fc1-activation-stage.wgsl.jinja": "Hx7fYjr8CwB2jXeTnh6G6IyavWrLwyOdtUsJwVEpgGk=",
17
  "qmoe-mix-stage.wgsl.jinja": "+deuFyuBn04UeIX8Hk2/ntCHNfM1vgkNDTcsh3uXVdI=",
18
  "qmoe-output-gemv.wgsl.jinja": "OJLoDDWMYE5yaTeDGY7qUeuf6AepJL7ZZLx8ko6OTEQ=",
19
- "qmoe-output-grouped-sgmat.wgsl.jinja": "XaFjSitK6tkK/QkrQ6wRConJLNekx+4u/FpxuFmh7DU=",
20
  "qmoe-output-grouped.wgsl.jinja": "mHkrsubKwd0YWYo4TMpl4HWdbWheoeiAsUU2WsUON8E=",
21
  "qmoe-output-stage.wgsl.jinja": "f9CWE9XR/5RTULNXdTHwA4gFsOL1QRg0H/rT3dD32Gs=",
22
  "qmoe-output-zero-inter.wgsl.jinja": "6iBjk5YqxSSoXufDo/aJJRS4v8ZOp7wF2aRRXU4m1SA=",
@@ -24,6 +24,6 @@
24
  "test.json": "3i3FSregL3rNfwD/Wi7X2Ix296NNB1zIV+Ap9TrpMeU="
25
  }
26
  },
27
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
28
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.QMoE" }
29
  }
 
1
  {
2
  "name": "com.microsoft.QMoE",
3
+ "id": "_com_microsoft_qmoe_webgpu_6caed83",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
9
  "files": {
10
  "bench.json": "zBVupnayA0Ggc4y737P9tzqyM/q/N8F4j+FEJXz6Wks=",
11
  "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
12
+ "manifest.json": "F9s2iIz0nNI2vJ+2zAFr55u8jaKjAEe5qvfmGxSc6mE=",
13
  "qmoe-fc1-activation-gemv.wgsl.jinja": "UvP8Vyx1E9+/+I09w/z9iB6h6MHrSrOmjKfIelyIS0M=",
14
+ "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "r/vlFM8WL1WfjkX7P2/DjKxrMDRKBmI+PtV/qQVLgCc=",
15
  "qmoe-fc1-activation-grouped.wgsl.jinja": "y2vGeNeXE5YmNrM2WrXogyF/fTvQ/nqYR4uWBTDvpPY=",
16
  "qmoe-fc1-activation-stage.wgsl.jinja": "Hx7fYjr8CwB2jXeTnh6G6IyavWrLwyOdtUsJwVEpgGk=",
17
  "qmoe-mix-stage.wgsl.jinja": "+deuFyuBn04UeIX8Hk2/ntCHNfM1vgkNDTcsh3uXVdI=",
18
  "qmoe-output-gemv.wgsl.jinja": "OJLoDDWMYE5yaTeDGY7qUeuf6AepJL7ZZLx8ko6OTEQ=",
19
+ "qmoe-output-grouped-sgmat.wgsl.jinja": "aoBOyGPOM2MXOpjvVCppacvWra4bVBewQEbjL2MPe2k=",
20
  "qmoe-output-grouped.wgsl.jinja": "mHkrsubKwd0YWYo4TMpl4HWdbWheoeiAsUU2WsUON8E=",
21
  "qmoe-output-stage.wgsl.jinja": "f9CWE9XR/5RTULNXdTHwA4gFsOL1QRg0H/rT3dD32Gs=",
22
  "qmoe-output-zero-inter.wgsl.jinja": "6iBjk5YqxSSoXufDo/aJJRS4v8ZOp7wF2aRRXU4m1SA=",
 
24
  "test.json": "3i3FSregL3rNfwD/Wi7X2Ix296NNB1zIV+Ap9TrpMeU="
25
  }
26
  },
27
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
28
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.QMoE" }
29
  }
build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja CHANGED
@@ -1,4 +1,7 @@
1
  enable subgroups;
 
 
 
2
  enable chromium_experimental_subgroup_matrix;
3
  diagnostic(off, chromium.subgroup_matrix_uniformity);
4
 
@@ -124,7 +127,7 @@ fn loadSHMB(expert: u32, weight_base: u32, k_idx: u32, row: u32, c_idx: u32) {
124
  }
125
  }
126
 
127
- @compute @workgroup_size({{ sgmatWorkgroup }}, 1, 1)
128
  fn main(
129
  @builtin(workgroup_id) wid: vec3<u32>,
130
  @builtin(local_invocation_index) local_idx: u32,
@@ -162,14 +165,14 @@ fn main(
162
 
163
  for (var step = 0u; step < TILE_K; step = step + 8u) {
164
  {% set operandScalar = "f32" %} let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
165
- var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
166
- var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
167
 
168
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
169
- var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
170
- var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
171
- var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
172
- var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
173
 
174
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
175
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -194,7 +197,7 @@ fn main(
194
  // The scratch banks are shared by both row groups, so the readback is fenced on both
195
  // sides of each collective store.
196
  {% for c in [0, 1, 2, 3] %}
197
- subgroupMatrixStore(&scratch[subtile_id][{{ c }}], 0u, matC{{ group }}{{ c }}, false, 8u);
198
  {% endfor %}
199
  workgroupBarrier();
200
  {
 
1
  enable subgroups;
2
+ {% if pinSubgroupSize32 %}
3
+ enable subgroup_size_control;
4
+ {% endif %}
5
  enable chromium_experimental_subgroup_matrix;
6
  diagnostic(off, chromium.subgroup_matrix_uniformity);
7
 
 
127
  }
128
  }
129
 
130
+ @compute @workgroup_size({{ sgmatWorkgroup }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
131
  fn main(
132
  @builtin(workgroup_id) wid: vec3<u32>,
133
  @builtin(local_invocation_index) local_idx: u32,
 
165
 
166
  for (var step = 0u; step < TILE_K; step = step + 8u) {
167
  {% set operandScalar = "f32" %} let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
168
+ var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
169
+ var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
170
 
171
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
172
+ var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
173
+ var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
174
+ var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
175
+ var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
176
 
177
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
178
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
 
197
  // The scratch banks are shared by both row groups, so the readback is fenced on both
198
  // sides of each collective store.
199
  {% for c in [0, 1, 2, 3] %}
200
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ c }}], 0u, matC{{ group }}{{ c }}, 8u);
201
  {% endfor %}
202
  workgroupBarrier();
203
  {
build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja CHANGED
@@ -1,4 +1,7 @@
1
  enable subgroups;
 
 
 
2
  enable chromium_experimental_subgroup_matrix;
3
  diagnostic(off, chromium.subgroup_matrix_uniformity);
4
 
@@ -118,7 +121,7 @@ fn storeOutput(col_base: u32, row: u32, col: u32, src_slot: u32, row_base: u32,
118
  }
119
  }
120
 
121
- @compute @workgroup_size({{ sgmatWorkgroup }}, 1, 1)
122
  fn main(
123
  @builtin(workgroup_id) wid: vec3<u32>,
124
  @builtin(local_invocation_index) local_idx: u32,
@@ -156,14 +159,14 @@ fn main(
156
 
157
  for (var step = 0u; step < TILE_K; step = step + 8u) {
158
  {% set operandScalar = "f32" %} let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
159
- var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
160
- var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset + 8u * TILE_K, false, TILE_K);
161
 
162
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
163
- var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
164
- var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 8u * TILE_K, true, TILE_K);
165
- var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 16u * TILE_K, true, TILE_K);
166
- var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset + 24u * TILE_K, true, TILE_K);
167
 
168
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
169
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -180,7 +183,7 @@ fn main(
180
  // Each scratch bank is written by a collective store and then read across lanes, and
181
  // the two row groups reuse the same banks, so both barriers are load-bearing.
182
  {% for col in [0, 1, 2, 3] %}
183
- subgroupMatrixStore(&scratch[subtile_id][{{ col }}], 0u, matC0{{ col }}, false, 8u);
184
  {% endfor %}
185
  workgroupBarrier();
186
  let row = sg_id / 4u;
@@ -189,7 +192,7 @@ fn main(
189
  workgroupBarrier();
190
 
191
  {% for col in [0, 1, 2, 3] %}
192
- subgroupMatrixStore(&scratch[subtile_id][{{ col }}], 0u, matC1{{ col }}, false, 8u);
193
  {% endfor %}
194
  workgroupBarrier();
195
  storeOutput(col_base + base_B, row, col, subtile_id, base_A + 8u, rows);
 
1
  enable subgroups;
2
+ {% if pinSubgroupSize32 %}
3
+ enable subgroup_size_control;
4
+ {% endif %}
5
  enable chromium_experimental_subgroup_matrix;
6
  diagnostic(off, chromium.subgroup_matrix_uniformity);
7
 
 
121
  }
122
  }
123
 
124
+ @compute @workgroup_size({{ sgmatWorkgroup }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
125
  fn main(
126
  @builtin(workgroup_id) wid: vec3<u32>,
127
  @builtin(local_invocation_index) local_idx: u32,
 
159
 
160
  for (var step = 0u; step < TILE_K; step = step + 8u) {
161
  {% set operandScalar = "f32" %} let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
162
+ var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
163
+ var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
164
 
165
  let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
166
+ var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
167
+ var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
168
+ var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
169
+ var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
170
 
171
  matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
172
  matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
 
183
  // Each scratch bank is written by a collective store and then read across lanes, and
184
  // the two row groups reuse the same banks, so both barriers are load-bearing.
185
  {% for col in [0, 1, 2, 3] %}
186
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ col }}], 0u, matC0{{ col }}, 8u);
187
  {% endfor %}
188
  workgroupBarrier();
189
  let row = sg_id / 4u;
 
192
  workgroupBarrier();
193
 
194
  {% for col in [0, 1, 2, 3] %}
195
+ subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ col }}], 0u, matC1{{ col }}, 8u);
196
  {% endfor %}
197
  workgroupBarrier();
198
  storeOutput(col_base + base_B, row, col, subtile_id, base_A + 8u, rows);