[ExecuTorch][WebGPU] Add et_vk.embedding_q4gsw (4-bit groupwise-symmetric quantized embedding)

[JulianCloudNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] et_vk.prepack test suite (export + native golden) #20292
• [ExecuTorch][WebGPU] Add et_vk.prepack (constant-tensor packing) for E2E weight loading #20265
• [ExecuTorch][WebGPU] et_vk.apply_rotary_emb test suite (export + native golden) #20290
• [ExecuTorch][WebGPU] Add et_vk.apply_rotary_emb (interleaved RoPE) + ValueList multi-output #20264
• [ExecuTorch][WebGPU] et_vk.embedding_q4gsw test suite (export + native golden) #20289
• -> [ExecuTorch][WebGPU] Add et_vk.embedding_q4gsw (4-bit groupwise-symmetric quantized embedding) #20263

Adds the WebGPU backend handler for et_vk.embedding_q4gsw.default (a 4-bit groupwise-symmetric quantized embedding gather) plus the host-side integer-input infra it requires.

The op is a single compute dispatch composed of one stage: one thread per 32-element block of each gathered row dequantizes the packed 4-bit table (q = (nibble - 8) * scale; even dim = high nibble, odd dim = low) into the fp32 output, mirroring the Vulkan embedding_q4gsw reference (flat buffer-backed weight; is_linear_weight=true is unsupported and throws). The workgroup size is a wg_size pipeline-override constant clamped to the device limit via WebGPUUtils::clamp_workgroup_size, the 1D dispatch count goes through WebGPUUtils::compute_1d_workgroup_count (validated before any GPU-object allocation), and the embedded WGSL string header is generated by gen_wgsl_headers.py.

Embedding indices arrive as int64 at the program boundary but the serialized graph stores them as int32, so the shared input path is extended with a host-side InputData view ({data, nbytes, host_is_int64}) and copy_inputs gains three branches: a byte-for-byte fast path when host and GPU sizes match, an int64->int32 narrowing copy when the buffer is int32 and the host input is twice as wide (mirrors the Vulkan kLong->kInt staging cast), and a fail-loud throw otherwise. WebGPUTensor gains elem_size/is_int to drive the narrowing decision, and update_symints_from_inputs takes the same InputData vector so execute() builds a single input list consumed by both.
@exported-using-ghexport

Differential Revision: D108428753

Differential Revision: D108428753

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20263

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 3 New Failures, 8 Pending, 1 Unrelated Failure

As of commit 7dea4c4 with merge base 0e65ba6 ():

NEW FAILURES - The following jobs have failed:

• Build Presets / apple (llm) / build (gh)
  The process '/opt/homebrew/bin/git' failed with exit code 1
• pull / unittest / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_mv3_model
• pull / unittest-editable / macos / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh) (trunk failure)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[JulianCloudNTH]

@claude review

[claude]

Claude finished @JulianCloudNTH's task in 2m 1s —— View job

---

Code Review: et_vk.embedding_q4gsw (WebGPU)

• Read PR diff and changed files
• Review op implementation (EmbeddingQ4gsw.cpp + WGSL)
• Review host-side input infra (WebGPUGraph, copy_inputs, InputData)
• Review tests
• Post review feedback

Overall this is a clean, well-validated addition. The op does extensive up-front shape/byte guards before allocating any GPU object, the WGSL nibble unpacking matches the documented pack convention, and the int64→int32 narrowing path is well reasoned. One correctness bug and a few smaller notes below.

---

🔴 copy_inputs can throw across the backend boundary

WebGPUBackend::execute deliberately wraps the symint refresh in a try/catch with the comment "Fail loud as a runtime Error so a throw never crosses the backend boundary" — but copy_inputs is called outside that block:

graph->copy_inputs(inputs);          // WebGPUBackend.cpp:109  — NOT guarded

try {
  graph->update_symints_from_inputs(inputs);   // :112-118 — guarded
  graph->propagate_resize();
} catch (const std::exception& e) { ... return Error::Internal; }

This PR newly makes copy_inputs throw std::runtime_error (the "fail-loud" branch in WebGPUGraph.cpp). Before this PR copy_inputs was non-throwing, so the unguarded call was fine — now an unsupported/mismatched input will let a C++ exception escape execute(), which returns Error and is not expected to throw. That contradicts the design comment two lines below and can crash the runtime instead of returning a clean error.

Suggest moving the copy_inputs call inside the existing try/catch (and the input-gathering loop is fine where it is):

try {
  graph->copy_inputs(inputs);
  graph->update_symints_from_inputs(inputs);
  graph->propagate_resize();
} catch (const std::exception& e) {
  ET_LOG(Error, "WebGPU input copy/symint refresh failed: %s", e.what());
  return Error::Internal;
}

Fix this →

---

🟡 Out-of-range indices read out of bounds (silent wrong output)

In the WGSL, token = u32(t_indices[indices_idx]) is used directly to index t_weight and t_scales with no clamp against vocab size. A token ≥ vocab (or a negative int32 that wraps to a huge u32) indexes past the weight/scales buffers. WebGPU bounds-checking will return 0 rather than crash, so the failure is silent. This mirrors the Vulkan reference's trust in valid indices, so it may be acceptable — but worth a one-line comment noting indices are assumed in-range, since unlike a CPU gather there's no bounds error to surface the mistake.

🟡 total_blocks narrowed to uint32 without a guard

total_blocks is computed as uint64_t but stored into the uint32_t EmbeddingParams::total_blocks and passed to compute_1d_workgroup_count via a static_cast<uint32_t>. There's no guard that it fits in 32 bits. All current test shapes are tiny (256 blocks), so this is theoretical, but a > UINT32_MAX check alongside the other guards would keep the fail-loud invariant consistent for large gathers.

🟢 Minor: per-execute allocation in the narrowing path

copy_inputs allocates a fresh std::vector<int32_t> narrowed(numel) on every execute for the int64→int32 case. Fine for correctness; if the embedding-gather is on a hot path you may later want a reusable scratch buffer. Not blocking.

---

Things that look good

• EmbeddingParams static_assert(sizeof == 32) + explicit _pad matches std140 uniform alignment — nice.
• Validation happens before any wgpuDeviceCreate* call, so failures don't leak GPU objects.
• The Python test computes the golden two independent ways (et_vk reference vs. a hand-rolled dequant+F.embedding) in test_golden_matches_eager, so a buggy golden can't mask a buggy kernel — good defensive testing.
• embedding_q4gsw_wgsl.h carries the wgsl-sha256 and DO NOT EDIT banner; assuming gen_wgsl_headers.py regenerated it, the embedded copy matches the .wgsl source (I diffed them — identical).
• The native C++ test exercises the int64→int32 narrowing path end-to-end by feeding an int64 tensor at the boundary.

Question

• The native test (test_webgpu_native.cpp) only runs the small shape (embed=64). The export helper supports llama1b (embed=2048, group=64) and the Python test_golden_matches_eager covers it, but the native GPU kernel is never run at that shape. Was that intentional (keep CI light), or worth adding so the per-group scale indexing is validated on-device too?

[JulianCloudNTH]

@claude review

[claude]

Claude finished @JulianCloudNTH's task in 3m 45s —— View job

---

Code Review: et_vk.embedding_q4gsw (WebGPU)

• Read PR diff and changed files
• Review op implementation (EmbeddingQ4gsw.cpp + WGSL)
• Verify dequant convention against the Vulkan reference
• Review host-side input infra (WebGPUGraph, copy_inputs, InputData)
• Re-check the prior review's findings against current HEAD

This is a clean, well-guarded addition. I re-verified the dequant math against the Vulkan flat-buffer reference and it matches exactly; the up-front validation discipline is good. Notes below, including which items from the earlier review are now resolved.

---

✅ Resolved since the last review

• copy_inputs now inside the try/catch. WebGPUBackend.cpp:110-117 moves graph->copy_inputs(inputs) into the guarded block, so the new fail-loud std::runtime_error branch in WebGPUGraph::copy_inputs is converted to Error::Internal and never crosses the backend boundary. The previous 🔴 is fixed.
• Out-of-range index behavior is now documented. embedding_q4gsw.wgsl:30 carries the // token assumed in-range (mirrors Vulkan; no vocab clamp) comment, which is what the earlier 🟡 asked for. Still a silent-zero on OOB (WebGPU bounds-checks to 0 rather than faulting), but it now matches the Vulkan contract explicitly.

✔️ Dequant convention verified against Vulkan

I diffed the unpacking against backends/vulkan/runtime/graph/ops/glsl/embedding_q4gsw.glsl (flat-buffer / non-linear path):

• Byte→nibble: even dim = high nibble (b >> 4), odd dim = low nibble (b & 0xF) — matches Vulkan's b0>>4 / b0&0xF pairing.
• q = (nibble - 8) * scale matches the +8-shifted pack.
• row_byte_base = token * (embed_dim/2) lines up with Vulkan's embedding_idx * blocks_per_row of 16-byte blocks, and little-endian byte addressing (word[byte_idx>>2], shift (byte_idx&3)*8) reads the same byte sequence as Vulkan's ivec4 indexing.
• Scale index token * groups_per_row + dim / group_size matches Vulkan exactly.

The embedded embedding_q4gsw_wgsl.h is byte-identical to the .wgsl source (sha256 banner present).

---

🟡 total_blocks truncated to uint32 without a guard (still open)

EmbeddingQ4gsw.cpp:110-111 computes total_blocks as uint64_t, then :139 passes static_cast<uint32_t>(total_blocks) to compute_1d_workgroup_count and :148 stores it into the uint32_t field. If total_blocks > UINT32_MAX the cast silently truncates before the dispatch-limit check, so the fail-loud invariant the rest of this function maintains is bypassed for very large gathers. All current shapes are tiny, so it's theoretical — but a one-line > UINT32_MAX guard alongside the other byte/shape checks would keep it consistent.

Fix this →

🟢 Per-execute allocation in the narrowing path (non-blocking)

WebGPUGraph.cpp:522 allocates a fresh std::vector<int32_t> narrowed(numel) on every execute for the int64→int32 case. Fine for correctness; if embedding-gather lands on a hot path, a reusable scratch buffer would avoid the per-call allocation. Not blocking.

🟢 Minor: Params.num_indices is unused in the WGSL

The uniform carries num_indices (set at EmbeddingQ4gsw.cpp:144) but the shader never reads it — the kernel derives indices_idx from block / blocks_per_row instead. Harmless (it's part of the 32-byte std140 layout), just dead in the shader. Worth a comment or dropping into _pad if you want the struct to reflect only what's consumed.

---

Note on tests

This PR's diff (7 files) contains no test files — the test_webgpu_native.cpp / test_golden_matches_eager coverage referenced earlier appears to live elsewhere in the ghstack (#20264 / #20289 / #20290), not here. Flagging only so the op + tests landing separately is intentional and not an omission. If the native test only exercises the small shape (embed=64), it's still worth running the llama1b shape (embed=2048, group=64) on-device at least once so the per-group scale indexing is validated on GPU and not just in the Python golden.

Looks good

• static_assert(sizeof(EmbeddingParams) == 32) + explicit _pad for std140 alignment.
• All shape/byte validation runs before any wgpuDeviceCreate*, so failures don't leak GPU objects.
• wg_size clamped to the device limit and the 1D dispatch count validated against maxComputeWorkgroupsPerDimension before allocation.
• The int64→int32 narrowing path is well-reasoned: byte-for-byte fast path, exact nbytes == tensor.nbytes * 2 width check for the cast, and a fail-loud throw otherwise.
  • Branch: gh/JulianCloudNTH/25/head