Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp

@@ -131,6 +131,10 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         "WebGPU embedding_q4gsw: dtype/byte-size mismatch "
         "(indices int32, weight uint8, scales/out fp32)");
   }
+  if (total_blocks > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU embedding_q4gsw: total_blocks exceeds uint32 dispatch range");
+  }
 
   // 1D dispatch: one thread per 32-dim block; validate before any alloc.
   const uint32_t wg_size =
@@ -141,7 +145,7 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   EmbeddingParams params = {};
   params.embed_dim = embed_dim;
   params.blocks_per_row = blocks_per_row;
-  params.num_indices = num_indices;
+  params.num_indices = num_indices; // std140 layout only; shader derives it
   params.group_size = static_cast<uint32_t>(group_size);
   params.groups_per_row = groups_per_row;
   params.bytes_per_row = bytes_per_row;

backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp

@@ -179,6 +179,11 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         "freqs shape != [seq, head_dim/2]");
   }
 
+  if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU apply_rotary_emb: pair count exceeds uint32 dispatch range");
+  }
+
   const uint32_t wg_size =
       utils::clamp_workgroup_size(device, kRotaryEmbeddingWorkgroupSizeX);
   // Validate both dispatches before any GPU-object alloc (no leak on throw).

backends/webgpu/scripts/test_webgpu_native_ci.sh

@@ -69,7 +69,7 @@ $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.embedding_q4gsw.test_embedding_q4gsw import export_embedding_q4gsw_model
 export_embedding_q4gsw_model('${EMBEDDING_MODEL}', '${EMBEDDING_GOLDEN}', '${EMBEDDING_INDICES}')
 export_embedding_q4gsw_model('${EMBEDDING_LLAMA1B_MODEL}', '${EMBEDDING_LLAMA1B_GOLDEN}', '${EMBEDDING_LLAMA1B_INDICES}', 'llama1b')
-" || echo "WARN: embedding_q4gsw export failed; webgpu_native_test embedding cases self-skip"
+" || echo "WARN: embedding_q4gsw export failed; embedding configs will FAIL in webgpu_native_test"
 
 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.rms_norm.test_rms_norm import export_rms_norm_cases

backends/webgpu/test/ops/embedding_q4gsw/test_embedding_q4gsw.py

@@ -122,14 +122,15 @@ def test_golden_matches_eager(self) -> None:
                 weight, scales, group_size = _quant_params(qm)
                 vocab = weight.shape[0]
                 embed = weight.shape[1] * 2
-                # fp64 reference dequant so the oracle carries no fp32 rounding.
-                deq = torch.empty((vocab, embed), dtype=torch.float64)
-                for row in range(vocab):
-                    for d in range(embed):
-                        byte = int(weight[row, d // 2])
-                        nib = (byte >> 4) if (d % 2 == 0) else (byte & 0xF)
-                        scale = float(scales[row, d // group_size])
-                        deq[row, d] = (nib - 8) * scale
+                # fp64 reference dequant, vectorized (no fp32 rounding in oracle).
+                w = weight.to(torch.int64)
+                nib = torch.empty((vocab, embed), dtype=torch.int64)
+                nib[:, 0::2] = (w >> 4) & 0xF  # even dim -> high nibble
+                nib[:, 1::2] = w & 0xF  # odd dim -> low nibble
+                scale_exp = scales.to(torch.float64).repeat_interleave(
+                    group_size, dim=1
+                )
+                deq = (nib - 8).to(torch.float64) * scale_exp
                 eager = torch.nn.functional.embedding(idx, deq)
                 golden = _golden(qm, idx)
                 torch.testing.assert_close(golden.double(), eager, atol=1e-5, rtol=1e-5)

backends/webgpu/test/test_webgpu_native.cpp

@@ -407,7 +407,16 @@ static float q4gsw_ramp(int i) {
   return static_cast<float>((i % 17) - 8) / 16.0f;
 }
 
-// Per-element dual tolerance (abs OR rel), parameterized like sdpa_within_tol.
+// Fwd decl of the per-element abs-OR-rel tolerance helper (defined below).
+static bool quant_within_tol(
+    const float* out,
+    const float* golden,
+    int n,
+    float atol,
+    float rtol,
+    float* ma,
+    float* mr);
+
 static std::vector<int32_t> load_indices(
     const std::string& path,
     size_t numel) {
@@ -483,19 +492,21 @@ static bool test_embedding_q4gsw(
   const float* out_data = out_tensor.const_data_ptr<float>();
 
   float max_abs_err = 0.0f, max_rel_err = 0.0f;
-  for (int i = 0; i < out_numel; i++) {
-    const float ae = std::abs(out_data[i] - golden[i]);
-    max_abs_err = std::max(max_abs_err, ae);
-    max_rel_err =
-        std::max(max_rel_err, ae / std::max(std::abs(golden[i]), 1e-6f));
-  }
+  const bool pass = quant_within_tol(
+      out_data,
+      golden.data(),
+      out_numel,
+      1e-3f,
+      1e-3f,
+      &max_abs_err,
+      &max_rel_err);
   printf(
       "Max abs error: %e   Max rel error: %e (checked %d elements)\n",
       max_abs_err,
       max_rel_err,
       out_numel);
-  if (max_abs_err > 1e-3f || max_rel_err > 1e-3f) {
-    printf("FAIL: embedding_q4gsw exceeds tolerance 1e-3\n");
+  if (!pass) {
+    printf("FAIL: embedding_q4gsw exceeds tolerance 1e-3 (abs AND rel)\n");
     return false;
   }
   printf("PASS: embedding_q4gsw test\n");