Update

[ghstack-poisoned]

Author: JCNTH

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -106,14 +106,13 @@ Error WebGPUBackend::execute(
         tensor.scalar_type() == executorch::aten::ScalarType::Long;
     inputs.push_back({tensor.const_data_ptr(), tensor.nbytes(), host_is_int64});
   }
-  graph->copy_inputs(inputs);
-
   // Fail loud as a runtime Error so a throw never crosses the backend boundary.
   try {
+    graph->copy_inputs(inputs);
     graph->update_symints_from_inputs(inputs);
     graph->propagate_resize();
   } catch (const std::exception& e) {
-    ET_LOG(Error, "WebGPU symint refresh/resize failed: %s", e.what());
+    ET_LOG(Error, "WebGPU input copy / symint refresh failed: %s", e.what());
     return Error::Internal;
   }
 

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -319,6 +319,10 @@ void WebGPUGraph::build(
                   "WebGPU: constant_id set but the constants table is missing "
                   "or the id is out of range");
             }
+          } else if (constant_id >= 0 && tensor.nbytes > 0) {
+            // constant_id set but constant_data null -> fail loud.
+            throw std::runtime_error(
+                "WebGPU: constant_id set but constant_data is null");
           }
         } else {
           // Shared buffer: track required size, defer allocation to pass 2

backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw.wgsl

@@ -27,6 +27,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let indices_idx = block / params.blocks_per_row;
   let base_dim = (block % params.blocks_per_row) * 32u;
 
+  // token assumed in-range (mirrors Vulkan; no vocab clamp).
   let token = u32(t_indices[indices_idx]);
   let row_byte_base = token * params.bytes_per_row;
   let out_base = indices_idx * params.embed_dim + base_dim;

backends/webgpu/runtime/ops/embedding_q4gsw/embedding_q4gsw_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from embedding_q4gsw.wgsl - DO NOT EDIT.
-// wgsl-sha256: b86f112c08bd73bdcbfed41a27fb63465ec36d1a50d0cdf2612684e30fceeac9
+// wgsl-sha256: 1fec9ed315696a88bb7db6c16454fc80e08ff73b0e39720b54515fda4ee1ef7c
 inline constexpr const char* kEmbeddingQ4gswWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_indices: array<i32>;
@@ -44,6 +44,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let indices_idx = block / params.blocks_per_row;
   let base_dim = (block % params.blocks_per_row) * 32u;
 
+  // token assumed in-range (mirrors Vulkan; no vocab clamp).
   let token = u32(t_indices[indices_idx]);
   let row_byte_base = token * params.bytes_per_row;
   let out_base = indices_idx * params.embed_dim + base_dim;

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

@@ -57,6 +57,7 @@ void add_rope_dispatch(
     uint32_t head_dim,
     uint32_t workgroup_count) {
   const uint32_t half_dim = head_dim / 2u;
+  // out.dims == in.dims (asserted in impl), so this matches the caller's wgc.
   const uint32_t num_pairs = static_cast<uint32_t>(numel_of(out.dims) / 2u);
 
   RotaryParams params = {};

backends/webgpu/scripts/test_webgpu_native_ci.sh

@@ -50,9 +50,9 @@ UPDATE_CACHE_OK=1
 EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte"
 EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin"
 EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
-ROPE_MODEL="/tmp/webgpu_rope.pte"
-ROPE_XQ_GOLDEN="/tmp/webgpu_rope_xq_golden.bin"
-ROPE_XK_GOLDEN="/tmp/webgpu_rope_xk_golden.bin"
+EMBEDDING_LLAMA1B_MODEL="/tmp/webgpu_embedding_q4gsw_llama1b.pte"
+EMBEDDING_LLAMA1B_INDICES="/tmp/webgpu_embedding_q4gsw_llama1b_indices.bin"
+EMBEDDING_LLAMA1B_GOLDEN="/tmp/webgpu_embedding_q4gsw_llama1b_golden.bin"
 
 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.add.test_add import export_add_model, export_chained_add_model
@@ -68,12 +68,8 @@ export_all_quantized_linear_models('/tmp')
 $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}')
-" || echo "WARN: embedding_q4gsw export failed; webgpu_native_test embedding case self-skips"
-
-$PYTHON_EXECUTABLE -c "
-from executorch.backends.webgpu.test.ops.rope.test_rope import export_rope_model
-export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}')
-" || echo "WARN: rope export failed; webgpu_native_test apply_rotary_emb case self-skips"
+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"
 
 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.rms_norm.test_rms_norm import export_rms_norm_cases
@@ -168,9 +164,9 @@ if [[ -x "${BIN_DIR}/webgpu_native_test" && -f "${PTE_MODEL}" ]]; then
       WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL="${EMBEDDING_MODEL}" \
       WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES="${EMBEDDING_INDICES}" \
       WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN="${EMBEDDING_GOLDEN}" \
-      WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \
-      WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \
-      WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \
+      WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \
+      WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \
+      WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \
       "${BIN_DIR}/webgpu_native_test"
 else
   echo "(skipping webgpu_native_test: no exported .pte — needs the executorch python wheel)"

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

@@ -122,7 +122,8 @@ def test_golden_matches_eager(self) -> None:
                 weight, scales, group_size = _quant_params(qm)
                 vocab = weight.shape[0]
                 embed = weight.shape[1] * 2
-                deq = torch.empty((vocab, embed), dtype=torch.float32)
+                # 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])
@@ -131,7 +132,7 @@ def test_golden_matches_eager(self) -> None:
                         deq[row, d] = (nib - 8) * scale
                 eager = torch.nn.functional.embedding(idx, deq)
                 golden = _golden(qm, idx)
-                torch.testing.assert_close(golden, eager, atol=1e-5, rtol=1e-5)
+                torch.testing.assert_close(golden.double(), eager, atol=1e-5, rtol=1e-5)
 
 
 def export_embedding_q4gsw_model(