sync c928d21e6cc1

Browse files

Files changed (5) hide show

build/webgpu/manifest.json +8 -8
build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja +32 -25
build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja +12 -9
build/webgpu/metadata.json +6 -6
build/webgpu/test.json +50 -0

build/webgpu/manifest.json CHANGED Viewed

@@ -640,7 +640,11 @@
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "packedFeature": "device.wgslLanguageFeatures.has(\"packed_4x8_integer_dot_product\")",
     "kBlocksExpected": "ceilDiv(attrs.K, attrs.block_size)",
     "blobSizeExpected": "ceilDiv(attrs.block_size * attrs.bits, 8)",
@@ -760,10 +764,9 @@
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
-      "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize == 32", "device.adapterInfo.subgroupMaxSize == 32", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": false,
         "hasBias": false,
@@ -947,10 +950,9 @@
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
-      "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize == 32", "device.adapterInfo.subgroupMaxSize == 32", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": true,
         "hasBias": true,
@@ -1134,10 +1136,9 @@
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
-      "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize == 32", "device.adapterInfo.subgroupMaxSize == 32", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": true,
         "hasBias": false,
@@ -1321,10 +1322,9 @@
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
-      "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize == 32", "device.adapterInfo.subgroupMaxSize == 32", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": false,
         "hasBias": true,

   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
+    "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",
+    "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
+    "wave32Effective": "wave32Adapter or pinSubgroupSize32",
     "packedFeature": "device.wgslLanguageFeatures.has(\"packed_4x8_integer_dot_product\")",
     "kBlocksExpected": "ceilDiv(attrs.K, attrs.block_size)",
     "blobSizeExpected": "ceilDiv(attrs.block_size * attrs.bits, 8)",
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
+      "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": false,
         "hasBias": false,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
+      "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": true,
         "hasBias": true,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
+      "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": true,
         "hasBias": false,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "limits": { "maxComputeWorkgroupStorageSize": 16384 },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
+      "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "attrs.block_size % 16 == 0", "attrs.K % 32 == 0", "aRows >= 32", "wave32Effective", "dispatchN64 <= device.limits.maxComputeWorkgroupsPerDimension", "sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension", "portableWorkgroupFits", "sgmatWorkgroupFits"],
       "constants": {
         "hasZero": false,
         "hasBias": true,

build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja CHANGED Viewed

@@ -8,9 +8,12 @@ enable f16;
 {% set nCols = gemvNCols %}
 {% set groups = 2 if nCols == 8 else 1 %}
 {% set comps = ["x", "y", "z", "w"] %}
 const WG: u32 = {{ workgroupSize }}u;
 const N_COLS: u32 = {{ nCols }}u;
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
@@ -18,25 +21,11 @@ var<workgroup> partials{{ sfx }}: array<vec4<f32>, {{ workgroupSize }}>;
 {% endfor %}
 // {{ bits }}-bit unpack: q2 packs four codes per byte, q4 two nibbles, q8 one byte/weight.
-{% macro matmul_nbits_packed_code(fn="packed_weight", buffer="b", kBlocks="params.kBlocks", blobSize="params.blobSize", bits=4) %}
-fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
-{% if bits == 2 %}
-  let byte_index = offset / 4u;
-  let shift = (offset % 4u) * 2u;
-  let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + byte_index;
-  return ({{ buffer }}[packed_index] >> shift) & 3u;
-{% elif bits == 4 %}
-  let byte_index = offset / 2u;
-  let shift = (offset % 2u) * 4u;
-  let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + byte_index;
-  return ({{ buffer }}[packed_index] >> shift) & 15u;
-{% else %}
-  let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
-  return {{ buffer }}[packed_index] & 255u;
-{% endif %}
 }
-{%- endmacro %}
-{{- matmul_nbits_packed_code(bits=bits) }}
 // Per-(column, block) zero point. Symmetric weights use the baked default
 // ({{ defaultZero }}); asymmetric weights carry one widened u32 per (col, block)
@@ -68,23 +57,41 @@ fn main(
   var acc{{ sfx }} = vec4<f32>(0.0);
 {% endfor %}
-  for (var k = tid; k < params.K; k = k + WG) {
     let block = k / params.blockSize;
-    let offset = k % params.blockSize;
-    let av = f32(a[k]);
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% for c in range(4) %}
 {% set i = g * 4 + c %}
 {% if i == 0 %}
     let scale0 = f32(scales[(col_base + 0u) * params.kBlocks + block]);
-    let q0 = f32(packed_weight(col_base + 0u, block, offset)) - zero_point({% if hasZero %}col_base + 0u, block{% endif %});
-    acc{{ sfx }}.x = acc{{ sfx }}.x + av * q0 * scale0;
 {% else %}
     if (col_base + {{ i }}u < params.N) {
       let scale{{ i }} = f32(scales[(col_base + {{ i }}u) * params.kBlocks + block]);
-      let q{{ i }} = f32(packed_weight(col_base + {{ i }}u, block, offset)) - zero_point({% if hasZero %}col_base + {{ i }}u, block{% endif %});
-      acc{{ sfx }}.{{ comps[c] }} = acc{{ sfx }}.{{ comps[c] }} + av * q{{ i }} * scale{{ i }};
     }
 {% endif %}
 {% endfor %}

 {% set nCols = gemvNCols %}
 {% set groups = 2 if nCols == 8 else 1 %}
 {% set comps = ["x", "y", "z", "w"] %}
+{% set codesPerByte = 4 if bits == 2 else (2 if bits == 4 else 1) %}
+{% set codeMask = 3 if bits == 2 else (15 if bits == 4 else 255) %}
 const WG: u32 = {{ workgroupSize }}u;
 const N_COLS: u32 = {{ nCols }}u;
+const CODES: u32 = {{ codesPerByte }}u;
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% endfor %}
 // {{ bits }}-bit unpack: q2 packs four codes per byte, q4 two nibbles, q8 one byte/weight.
+// The whole stored byte, addressed by its index within the block rather than by
+// one of the codes inside it.
+fn packed_byte(n: u32, block: u32, byte_index: u32) -> u32 {
+  return b[(n * params.kBlocks + block) * params.blobSize + byte_index];
 }
 // Per-(column, block) zero point. Symmetric weights use the baked default
 // ({{ defaultZero }}); asymmetric weights carry one widened u32 per (col, block)
   var acc{{ sfx }} = vec4<f32>(0.0);
 {% endfor %}
+  // Stride whole bytes: every code a byte holds is consumed from one fetch.
+  for (var k = tid * CODES; k < params.K; k = k + WG * CODES) {
     let block = k / params.blockSize;
+    // blockSize is a power of two of at least 16, so a byte never straddles two
+    // blocks and every code below shares this block's scale and zero point.
+    let byte_index = (k % params.blockSize) / CODES;
+{% for j in range(codesPerByte) %}
+{% if j == 0 %}
+    let av0 = f32(a[k]);
+{% else %}
+    // K need not be a multiple of CODES; a code past the end contributes zero.
+    let av{{ j }} = select(0.0, f32(a[min(k + {{ j }}u, params.K - 1u)]), k + {{ j }}u < params.K);
+{% endif %}
+{% endfor %}
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% for c in range(4) %}
 {% set i = g * 4 + c %}
 {% if i == 0 %}
     let scale0 = f32(scales[(col_base + 0u) * params.kBlocks + block]);
+    let zero0 = zero_point({% if hasZero %}col_base + 0u, block{% endif %});
+    let word0 = packed_byte(col_base + 0u, block, byte_index);
+{% for j in range(codesPerByte) %}
+    acc{{ sfx }}.x = acc{{ sfx }}.x
+      + av{{ j }} * (f32((word0 >> {{ j * bits }}u) & {{ codeMask }}u) - zero0) * scale0;
+{% endfor %}
 {% else %}
     if (col_base + {{ i }}u < params.N) {
       let scale{{ i }} = f32(scales[(col_base + {{ i }}u) * params.kBlocks + block]);
+      let zero{{ i }} = zero_point({% if hasZero %}col_base + {{ i }}u, block{% endif %});
+      let word{{ i }} = packed_byte(col_base + {{ i }}u, block, byte_index);
+{% for j in range(codesPerByte) %}
+      acc{{ sfx }}.{{ comps[c] }} = acc{{ sfx }}.{{ comps[c] }}
+        + av{{ j }} * (f32((word{{ i }} >> {{ j * bits }}u) & {{ codeMask }}u) - zero{{ i }}) * scale{{ i }};
+{% endfor %}
     }
 {% endif %}
 {% endfor %}

build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja CHANGED Viewed

@@ -8,6 +8,9 @@
 // Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
 // sees partial 8x8 tiles; M tails handled by row_limit + the loadSHMA zero-fill.
 enable subgroups;
 enable chromium_experimental_subgroup_matrix;
 {% if usesF16 %}
 enable f16;
@@ -118,7 +121,7 @@ fn storeOutput(offset: u32, col_base: u32, row: u32, col: u32, src_slot: u32, ro
   }
 }
-@compute @workgroup_size({{ workgroupSize }}, 1, 1)
 fn main(
   @builtin(workgroup_id) workgroup_id: vec3<u32>,
   @builtin(local_invocation_index) local_idx: u32,
@@ -145,14 +148,14 @@ fn main(
     for (var step = 0u; step < TILE_K; step = step + 8u) {
 {% set operandScalar = "f32" %}      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>>(&tile_A, matrix_a_offset, false, TILE_K);
-      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);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>>(&tile_B, matrix_b_offset, true, TILE_K);
-      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);
-      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);
-      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);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
@@ -172,7 +175,7 @@ fn main(
   // clobbering the first's still-in-flight readback when a partial final M-tile
   // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
 {% for col in [0, 1, 2, 3] %}
-  subgroupMatrixStore(&scratch[subtile_id][{{ col }}], 0u, matC0{{ col }}, false, 8u);
 {% endfor %}
   workgroupBarrier();
   let row = sg_id / 4u;
@@ -184,7 +187,7 @@ fn main(
   workgroupBarrier();
 {% for col in [0, 1, 2, 3] %}
-  subgroupMatrixStore(&scratch[subtile_id][{{ col }}], 0u, matC1{{ col }}, false, 8u);
 {% endfor %}
   workgroupBarrier();
   matrix_c_offset = matrix_c_offset + 8u * N;

 // Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
 // sees partial 8x8 tiles; M tails handled by row_limit + the loadSHMA zero-fill.
 enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
 enable chromium_experimental_subgroup_matrix;
 {% if usesF16 %}
 enable f16;
   }
 }
+@compute @workgroup_size({{ workgroupSize }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
 fn main(
   @builtin(workgroup_id) workgroup_id: vec3<u32>,
   @builtin(local_invocation_index) local_idx: u32,
     for (var step = 0u; step < TILE_K; step = step + 8u) {
 {% set operandScalar = "f32" %}      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
+      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);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
+      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);
+      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);
+      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);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);
   // clobbering the first's still-in-flight readback when a partial final M-tile
   // diverges storeOutput's guard. Without both barriers the last valid row can be corrupted.
 {% for col in [0, 1, 2, 3] %}
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ col }}], 0u, matC0{{ col }}, 8u);
 {% endfor %}
   workgroupBarrier();
   let row = sg_id / 4u;
   workgroupBarrier();
 {% for col in [0, 1, 2, 3] %}
+  subgroupMatrixStore<row_major>(&scratch[subtile_id][{{ col }}], 0u, matC1{{ col }}, 8u);
 {% endfor %}
   workgroupBarrier();
   matrix_c_offset = matrix_c_offset + 8u * N;

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.MatMulNBits",
-  "id": "_com_microsoft_matmulnbits_webgpu_81b817c",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,18 +8,18 @@
     "algorithm": "sha256",
     "files": {
       "bench.json": "BbpRtHBH/vkJDaybyRBNKhuUCpe4XMJtKmW3bhRcEZE=",
-      "manifest.json": "Gz92F8En4SeyVGrSZBtQghet4kVvjHxGIplFWSIq7AM=",
       "matmul-nbits-dp4a-quantize.wgsl.jinja": "WEUGSDH5qWYg10f/mkAHxq8AkIkMW9R/Wn/Y23Q0Jw0=",
-      "matmul-nbits-gemv-q4.wgsl.jinja": "r9eGhKkZPaNaUHFcDZh1JpS+w6z4Wow/7wGrLMVVgsA=",
       "matmul-nbits-q4-dp4a-prefill.wgsl.jinja": "qUrURZDpPzMpABnTgESkHgLCQiAWxEZ8wvzkAPTW2kI=",
       "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja": "HjZraSjOpaSl5Cjns4scselOIn+ibxFEu+9LTeUI0QI=",
       "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja": "oU7BcxJnKIEfgnkWHPq77RrCmrE0iIm6ED7cgVfZUYs=",
       "matmul-nbits-q4-prefill-tiled.wgsl.jinja": "TdetCiv5o88hvWCGNDewpFX+M+NZrz9rekU/d5kMZTM=",
-      "matmul-nbits-q4-sgmat.wgsl.jinja": "8E4yGRvezICz5ctdW4g2w/HBfLvqJF2hA+jOobcuJio=",
       "matmul-nbits.wgsl.jinja": "DRkwODEER11pYQxt2ssS9YvFggO3CbUqQ9yag53J5lY=",
-      "test.json": "2mpJ5fx2GnlZycPygUyHbOM37n7X6VZmUMGHNkJa7qw="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MatMulNBits" }
 }

 {
   "name": "com.microsoft.MatMulNBits",
+  "id": "_com_microsoft_matmulnbits_webgpu_5495bd2",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "BbpRtHBH/vkJDaybyRBNKhuUCpe4XMJtKmW3bhRcEZE=",
+      "manifest.json": "z8oCYWRaMHew5ogYuf7Txv4/6y2qpVoVODDlLxEZ+vA=",
       "matmul-nbits-dp4a-quantize.wgsl.jinja": "WEUGSDH5qWYg10f/mkAHxq8AkIkMW9R/Wn/Y23Q0Jw0=",
+      "matmul-nbits-gemv-q4.wgsl.jinja": "x1LxVx2cJBdN2cqapXz2T6AScpTcbmTO+B3tzvIXuT8=",
       "matmul-nbits-q4-dp4a-prefill.wgsl.jinja": "qUrURZDpPzMpABnTgESkHgLCQiAWxEZ8wvzkAPTW2kI=",
       "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja": "HjZraSjOpaSl5Cjns4scselOIn+ibxFEu+9LTeUI0QI=",
       "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja": "oU7BcxJnKIEfgnkWHPq77RrCmrE0iIm6ED7cgVfZUYs=",
       "matmul-nbits-q4-prefill-tiled.wgsl.jinja": "TdetCiv5o88hvWCGNDewpFX+M+NZrz9rekU/d5kMZTM=",
+      "matmul-nbits-q4-sgmat.wgsl.jinja": "PfO9QrUyNflp7RtUbUTRDwx7GDI6UUwz9I+6mJpO1OA=",
       "matmul-nbits.wgsl.jinja": "DRkwODEER11pYQxt2ssS9YvFggO3CbUqQ9yag53J5lY=",
+      "test.json": "XLwnbLdOJASke2yBnDScjJi+o2XISmdRCebb2jwWpDI="
     }
   },
+  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MatMulNBits" }
 }

build/webgpu/test.json CHANGED Viewed

@@ -858,6 +858,56 @@
       "outputs": { "yT": { "dtype": "float32", "shape": [1, 7], "tolerance": 0.00001 } },
       "attrs": { "K": 32, "N": 7, "bits": 4, "block_size": 32 }
     },
     {
       "name": "q4_gemv_default_zero_m1_n13_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },

       "outputs": { "yT": { "dtype": "float32", "shape": [1, 7], "tolerance": 0.00001 } },
       "attrs": { "K": 32, "N": 7, "bits": 4, "block_size": 32 }
     },
+    {
+      "name": "q4_gemv_default_zero_m1_odd_k33_n5",
+      "provenance": {
+        "notes": "K=33 is odd, so the last stored byte of the second block carries one live nibble and one past the end of K. The GEMV walks whole bytes, so this is the only shape that exercises its out-of-range activation guard; every other M==1 case has an even K."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [5, 2, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [5, 2],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.00001 } },
+      "attrs": { "K": 33, "N": 5, "bits": 4, "block_size": 32 }
+    },
+    {
+      "name": "q2_gemv_default_zero_m1_k18_n5",
+      "provenance": {
+        "notes": "The only two-bit M==1 case: four codes share one stored byte, and K=18 leaves two of the last byte's four codes past the end of K. Covers both the two-bit unpack in the GEMV and its out-of-range activation guard at the widest code-per-byte packing."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 18],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [5, 2, 4],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195, 84, 237] }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [5, 2],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.00001 } },
+      "attrs": { "K": 18, "N": 5, "bits": 2, "block_size": 16 }
+    },
     {
       "name": "q4_gemv_default_zero_m1_n13_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },