Skip to content

Commit 8200b9e

Browse files
JCNTHJCNTH
authored
Refresh backend README (pytorch#21385) (pytorch#21385)
Summary: **Document the current WebGPU backend and measured browser results** The existing README describes an early five-operator prototype. This refresh aligns it with the landed backend and validated browser artifacts. Key changes: - `README.md` — document the current architecture and representative 100+ operator registry - `README.md` — report accepted Llama browser medians and conservative Qwen validation status - `README.md` — replace stale setup and test guidance with current paths Documentation only; runtime behavior is unchanged. Co-authored-with: Claude Code. Reviewed By: psiddh Differential Revision: D113580415
1 parent 30bf6af commit 8200b9e

1 file changed

Lines changed: 158 additions & 97 deletions

File tree

‎backends/webgpu/README.md‎

Lines changed: 158 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -1,153 +1,214 @@
1-
# WebGPU Backend
1+
# ExecuTorch WebGPU Backend
22

3-
Run ExecuTorch models on the GPU via [WebGPU](https://www.w3.org/TR/webgpu/). The backend compiles delegated subgraphs into WGSL compute shaders executed natively through [Dawn](https://dawn.googlesource.com/dawn), whose Tint compiler is the reference WGSL implementation (Metal on macOS, Vulkan on Linux/Windows).
3+
Run ExecuTorch models on the GPU through
4+
[WebGPU](https://www.w3.org/TR/webgpu/). The backend compiles delegated
5+
subgraphs into WGSL compute shaders and executes them through
6+
[Dawn](https://dawn.googlesource.com/dawn) and its Tint WGSL compiler. The same
7+
graph execution design targets Metal and Vulkan in native builds. Browser
8+
builds compile the same execution path with Emscripten and `emdawnwebgpu`.
49

5-
> **Status: Prototype, under active development.** The backend runs the core of transformer inference today — `add`, `rms_norm`, fused scaled-dot-product attention with KV cache, and 4-bit weight-only quantized linear — plus quantized embedding, rotary embedding, and constant prepacking. See [Progress](#progress) for shipped milestones.
10+
> **Status: under active development.** The backend has more than 100
11+
> registered operator symbols and has demonstrated browser artifacts and
12+
> workflows across language, vision, retrieval, audio, and on-device training.
613
7-
## Progress
14+
## Performance
815

9-
Milestones landed on `main`:
16+
Measured in July 2026 on an Apple M4 Pro in Chrome Canary, using WebGPU
17+
GPU-timestamp queries and a hash-pinned, optimized dynamic-shape 4-bit Llama
18+
3.2 1B artifact. Headline values are medians from three measured runs per cell
19+
after three untimed prefill and decode warmups.
1020

11-
| Date | Milestone | Pull Request |
12-
|---|---|---|
13-
| 2026-04 | Made it possible to run ExecuTorch models on the GPU through WebGPU — built the backend from the ground up, including the runtime delegate that builds the GPU graph (buffers, pipelines, bind groups) and runs the model on Metal and Vulkan | [#18808](https://github.com/pytorch/executorch/pull/18808) |
14-
| 2026-06 | Grew model support beyond element-wise operators — added the root-mean-square normalization operator (`rms_norm`) and named-data weight loading | [#19963](https://github.com/pytorch/executorch/pull/19963) |
15-
| 2026-06 | Made sure every change is automatically tested — added WebGPU to ExecuTorch's standard backend test suite, running on Linux/x86 in CI | [#19964](https://github.com/pytorch/executorch/pull/19964) |
16-
| 2026-06 | Removed a class of bugs and manual upkeep — the WGSL shaders are now generated automatically, with a build-time check that fails the build on shader/source drift | [#19981](https://github.com/pytorch/executorch/pull/19981) |
17-
| 2026-06 | Got the test suite to actually run work on the GPU — added operator-allowlist delegation (unsupported operations fall back to the CPU) and a process-wide GPU device context, so models execute on the GPU during testing | [#20036](https://github.com/pytorch/executorch/pull/20036) |
18-
| 2026-06 | Made testing match the WebGPU standard exactly — switched the native runtime and tests to Google's Dawn shader compiler (Tint, the source-of-truth WGSL implementation) running on SwiftShader for headless GPU execution | [#20079](https://github.com/pytorch/executorch/pull/20079) |
19-
| 2026-06 | Strengthened correctness for models that run in several GPU passes — added dispatch-ordering and scratch-buffer (temporary GPU memory) support and tests | [#20080](https://github.com/pytorch/executorch/pull/20080) |
20-
| 2026-06 | Added the attention core of transformer inference — fused scaled-dot-product attention (`sdpa_with_kv_cache`) with an `update_cache` operator for autoregressive decode | [#20086](https://github.com/pytorch/executorch/pull/20086), [#20087](https://github.com/pytorch/executorch/pull/20087) |
21-
| 2026-06 | Added on-GPU kernel timing via WebGPU timestamp queries, for true GPU-side profiling | [#20201](https://github.com/pytorch/executorch/pull/20201) |
22-
| 2026-06 | Added the dominant compute in quantized LLMs — 4-bit weight-only quantized linear (`linear_q4gsw`), a dequantize-and-matmul kernel | [#20226](https://github.com/pytorch/executorch/pull/20226), [#20227](https://github.com/pytorch/executorch/pull/20227) |
21+
| Context / prompt tokens | Prefill | Decode |
22+
|---:|---:|---:|
23+
| 128 | 1,881.5 tokens/s | 188.3 tokens/s |
24+
| 512 | 2,291.0 tokens/s | 179.5 tokens/s |
25+
| 1,024 | 2,209.7 tokens/s | 167.2 tokens/s |
26+
| 2,048 | 1,950.7 tokens/s | 152.9 tokens/s |
27+
| 4,096 | 1,598.4 tokens/s | 132.8 tokens/s |
28+
| 8,192 | 1,164.2 tokens/s | 96.2 tokens/s |
2329

24-
In review:
30+
These figures describe this artifact and device, not the performance of every
31+
model or hardware configuration.
2532

26-
| Milestone | Pull Request |
33+
## Optimization Techniques
34+
35+
| Category | Optimizations |
36+
|---|---|
37+
| Prefill | Shared-memory tiled "Steel" GEMM with fp16 multiplication and fp32 accumulation, shape-routed kernel selection, scoped output suppression, and QKV fusion |
38+
| Decode | Cooperative GEMV with bi-column batching, FlashDecoding with split KV, and register-tiled SDPA |
39+
| SDPA | Register-tiled QK and AV, coalesced KV reads, causal-tile skipping, and aligned loads |
40+
| Quantization | Packed-word dequantization, 4-bit weight-only kernels, dynamic 8da4w quantization, and int8 q8ta kernels |
41+
| Memory | f16 KV cache and scratch-pool reuse |
42+
| Dispatch | Two-dimensional folded dispatch, runtime-configurable workgroup sizes, and SwiGLU fusion |
43+
44+
## Demonstrated Tasks
45+
46+
| Task | Capability |
2747
|---|---|
28-
| Adds 4-bit quantized embedding (`embedding_q4gsw`) | [#20263](https://github.com/pytorch/executorch/pull/20263) |
29-
| Adds rotary position embedding / RoPE (`apply_rotary_emb`) | [#20264](https://github.com/pytorch/executorch/pull/20264) |
30-
| Adds constant prepacking (`prepack`) for end-to-end model weight handling | [#20265](https://github.com/pytorch/executorch/pull/20265) |
48+
| On-device chat | Streaming local language-model generation |
49+
| Retrieval-augmented generation | Semantic retrieval and grounded generation over user documents |
50+
| Image segmentation | Prompted object grounding and segmentation |
51+
| Vision-language understanding | Image question answering and captioning |
52+
| Depth estimation | Per-pixel depth estimation from a single image |
53+
| Scene understanding | Scene classification and shoppable-product detection |
54+
| Object detection | Image object detection and labeling |
55+
| Super-resolution | 4x image upscaling |
56+
| Background removal | Portrait matting and background removal |
57+
| Speech-to-text | Microphone and uploaded-audio transcription |
58+
| On-device fine-tuning | GPU-resident adapter training followed by personalized generation |
59+
60+
These rows describe demonstrated artifacts and workflows. The registered
61+
operator surface is broader, but registry coverage alone does not guarantee
62+
that an arbitrary model will run end to end.
63+
64+
## Operator Support
65+
66+
The backend registers more than 100 operator symbols. Representative groups
67+
are listed below; this is not an exhaustive registry listing.
68+
69+
| Category | Representative operators |
70+
|---|---|
71+
| Quantized linear and embedding | `linear_q4gsw`, `embedding_q4gsw`, `linear_qcs4w`, `linear_dq8ca_q4gsw`, `choose_qparams_affine` |
72+
| Int8 quantized (q8ta) | `q8ta_linear`, `q8ta_conv2d`, `q8ta_add`, `q8ta_relu`, `q8ta_pixel_shuffle`, per-tensor quantize/dequantize |
73+
| Training primitives | `adamw_step`, `fused_ce`, `linear_q4gsw_backward`, `linear_dW`, `q4gsw_requant` |
74+
| Attention and position | `sdpa_with_kv_cache`, `update_cache`, `apply_rotary_emb` (default, interleaved, and Hugging Face layouts) |
75+
| Elementwise | `add`, `mul`, `sub`, `div`, `pow`, `minimum`, `sigmoid`, `relu`, `gelu`, `tanh`, `abs`, `neg`, `exp`, `sqrt`, `rsqrt`, `sin`, `cos`, `clamp` |
76+
| Comparison and boolean | `eq`, `ne`, `le`, `ge`, `lt`, `gt`, `where`, logical and bitwise operators |
77+
| Shape and memory | `view_copy`, `slice_copy`, `select_copy`, `cat`, `permute_copy`, `squeeze_copy`, `unsqueeze_copy`, `clone`, `expand_copy`, `fill`, `repeat`, `flip`, `pixel_shuffle` |
78+
| Reduction and normalization | `rms_norm`, `layer_norm`, `batch_norm`, `group_norm`, `softmax`, `log_softmax`, `sum`, `mean`, `argmax`, `argmin`, `amax`, `amin`, pooling |
79+
| Convolution and spatial | `conv1d`, `conv_with_clamp`, `grid_sampler_2d`, `grid_priors`, `upsample_bilinear2d` |
80+
| Matrix multiplication | `mm`, `bmm`, `linear` |
81+
| Indexing | `index.Tensor`, `gather`, `embedding`, `index_select` |
82+
| Infrastructure | `prepack`, `select_as_symint`, `_to_copy` |
3183

3284
## Architecture
3385

34-
```
86+
```text
3587
PyTorch model
3688
│ torch.export
3789
▼
3890
Exported Program
39-
│ VulkanPartitioner (tags supported fp32 ops)
91+
│ VulkanPartitioner (tags supported ops)
4092
▼
4193
Edge Dialect IR
42-
│ VulkanBackend.preprocess (builds Vulkan FlatBuffer, buffer-only storage)
94+
│ VulkanBackend.preprocess (builds a Vulkan FlatBuffer)
4395
▼
4496
.pte file (with VH00/VK00 delegate blob)
4597
│
46-
▼
47-
Native runtime (Dawn/Tint → Metal / Vulkan)
48-
│ WebGPUGraph::build → creates GPU buffers, pipelines, bind groups
49-
│ WebGPUGraph::execute → encodes + submits compute passes
50-
▼
51-
GPU output (mapped back to CPU via wgpuDevicePoll)
98+
├── Native runtime (Dawn/Tint → Metal or Vulkan)
99+
│ └── WebGPUGraph::build → creates GPU buffers, pipelines, bind groups
100+
│ WebGPUGraph::execute → encodes and submits compute passes
101+
│
102+
└── Browser runtime (Emscripten + emdawnwebgpu)
103+
└── Same graph execution path, compiled to WebAssembly
52104
```
53105

54106
Key design choices:
55-
- **Reuses Vulkan serialization** — the delegate blob is a Vulkan FlatBuffer (`VK00`) with a `VH00` header. All tensor storage is forced to `BUFFER` (WebGPU has no 3D storage textures).
56-
- **Built-in WGSL shaders** — shader source is compiled as C++ string constants. Future work will embed fused shaders in the FlatBuffer for compile-time mega-kernel fusion.
57-
- **No Python AOT code** — directly consumes .pte files exported via `VulkanPartitioner`.
58-
59-
## Operator Support
60107

61-
| Operator | WGSL Shader | Notes |
62-
|---|---|---|
63-
| `aten.add.Tensor` | `binary_add.wgsl` | Element-wise with alpha: `out = in1 + alpha * in2` |
64-
| `et_vk.rms_norm.default` | `rms_norm.wgsl` | Root-mean-square normalization |
65-
| `sdpa_with_kv_cache.default` | `sdpa_compute_attn_weights.wgsl`, `sdpa_softmax.wgsl`, `sdpa_compute_out.wgsl` | Fused scaled-dot-product attention (QK / softmax / AV) with KV cache |
66-
| `llama.update_cache.default` | `update_cache.wgsl` | In-place KV cache update for autoregressive decode |
67-
| `et_vk.linear_q4gsw.default` | `q4gsw_linear.wgsl` | 4-bit weight-only quantized linear (dequantize + matmul) |
68-
69-
**In review:** quantized embedding (`embedding_q4gsw`), rotary embedding (`apply_rotary_emb`), and constant prepacking (`prepack`).
70-
71-
**Planned:** `mul`, `sigmoid`, shape ops (`view`, `permute`, `slice`, `select`, `cat`, `squeeze`/`unsqueeze`), and `index` — the remaining ops needed for end-to-end Llama 3.2 1B.
72-
73-
## Quick Start
74-
75-
### 1. Setup
76-
77-
```bash
78-
bash backends/webgpu/scripts/setup-wgpu-native.sh
79-
```
80-
81-
This downloads prebuilt wgpu-native binaries for your platform.
82-
83-
### 2. Export a model
108+
- **Reuses Vulkan serialization.** The delegate blob is a Vulkan FlatBuffer
109+
(`VK00`) with a `VH00` header. The WebGPU runtime ignores Vulkan texture
110+
storage annotations and allocates tensors as buffers.
111+
- **Built-in WGSL shaders.** Shader source is embedded as C++ string constants,
112+
with build-time code generation and drift validation.
113+
- **No Python AOT layer.** The backend directly consumes `.pte` files exported
114+
with `VulkanPartitioner`.
115+
- **Dynamic shapes.** Tensors allocate at their maximum shape, with SymInt
116+
arithmetic and per-operator resize hooks for runtime dimensions.
117+
- **Shape-routed dispatch.** Runtime tensor dimensions select specialized
118+
kernels, such as tiled GEMM for prefill and cooperative GEMV for decode.
119+
120+
## Infrastructure and Testing
121+
122+
- **CI:** Dawn/Tint and SwiftShader provide headless GPU execution on Linux/x86.
123+
- **Operator tests:** Python operator-test modules exercise export and
124+
delegation. The code-generated op-test catalog and native runners compare
125+
GPU output with eager-generated goldens.
126+
- **GPU profiling:** WebGPU timestamp queries provide per-kernel GPU timings on
127+
devices that expose the feature.
128+
- **Shader code generation:** `gen_wgsl_headers.py` generates embedded
129+
`*_wgsl.h` headers; source/header drift fails validation.
130+
131+
## Linux Native Quick Start
132+
133+
The `test_build_webgpu.sh` flow sources
134+
`.ci/scripts/setup-webgpu-linux-deps.sh` to install Dawn and SwiftShader
135+
prebuilts on Linux. On macOS, provide a configured Dawn installation instead
136+
of using this script.
137+
138+
### 1. Export an illustrative model
84139

85140
```python
86141
import torch
87-
from executorch.backends.vulkan import VulkanPartitioner
142+
143+
from executorch.backends.vulkan.partitioner.vulkan_partitioner import (
144+
VulkanPartitioner,
145+
)
88146
from executorch.exir import to_edge_transform_and_lower
89147

90-
class AddModule(torch.nn.Module):
91-
def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
92-
return a + b
93148

94-
ep = torch.export.export(AddModule(), (torch.randn(4, 4), torch.randn(4, 4)))
149+
class AddOne(torch.nn.Module):
150+
def forward(self, x: torch.Tensor) -> torch.Tensor:
151+
return x + 1.0
152+
153+
154+
ep = torch.export.export(AddOne(), (torch.randn(4, 4),))
95155
et_program = to_edge_transform_and_lower(
96156
ep, partitioner=[VulkanPartitioner()]
97157
).to_executorch()
98158

99-
with open("add.pte", "wb") as f:
100-
f.write(et_program.buffer)
159+
with open("model.pte", "wb") as file:
160+
file.write(et_program.buffer)
101161
```
102162

103-
### 3. Build and run
163+
This snippet demonstrates the export path. The validation script below exports
164+
and runs its own native reference models; it does not consume `model.pte`.
165+
166+
### 2. Build and run native validation
104167

105168
```bash
106169
bash backends/webgpu/test/test_build_webgpu.sh
107170
```
108171

109-
This runs Python export tests, exports a .pte, builds the native runtime, and validates GPU output.
172+
The script exports a `.pte`, builds the native runtime, and validates GPU
173+
output.
110174

111175
## Directory Structure
112176

113-
```
177+
```text
114178
backends/webgpu/
115179
├── CMakeLists.txt
116180
├── README.md
117181
├── runtime/
118-
│ ├── WebGPUBackend.h/cpp # BackendInterface (init/execute)
119-
│ ├── WebGPUGraph.h/cpp # GPU graph: buffers, pipelines, dispatch
120-
│ ├── WebGPUDelegateHeader.h/cpp # VH00 header parser
121-
│ ├── WebGPUDevice.h/cpp # Dawn device abstraction
122-
│ ├── WebGPUUtils.h # Workgroup-size helpers
123-
│ └── ops/
124-
│ ├── OperatorRegistry.h/cpp # Op dispatch table
182+
│ ├── WebGPUBackend.h/cpp # BackendInterface (init/execute)
183+
│ ├── WebGPUGraph.h/cpp # GPU graph: buffers, pipelines, dispatch
184+
│ ├── WebGPUDelegateHeader.h/cpp # VH00 header parser
185+
│ ├── WebGPUDevice.h/cpp # Dawn device abstraction
186+
│ ├── WebGPUUtils.h # Workgroup-size helpers
187+
│ └── ops/ # Operator implementations
188+
│ ├── OperatorRegistry.h/cpp # Operator dispatch table
125189
│ ├── add/
126-
│ │ ├── BinaryOp.cpp # aten.add.Tensor implementation
127-
│ │ ├── binary_add.wgsl # WGSL shader source
128-
│ │ └── binary_add_wgsl.h # Shader as C++ string constant
129-
│ └── rms_norm/
130-
│ ├── RmsNorm.cpp # et_vk.rms_norm implementation
131-
│ ├── rms_norm.wgsl # WGSL shader source
132-
│ └── rms_norm_wgsl.h # Shader as C++ string constant
190+
│ │ ├── BinaryOp.cpp
191+
│ │ ├── binary_add.wgsl
192+
│ │ └── binary_add_wgsl.h
193+
│ └── ...
133194
├── scripts/
134-
│ ├── setup-wgpu-native.sh # Download wgpu-native binaries
135-
│ └── gen_wgsl_headers.py # Generate the embedded *_wgsl.h shader headers
195+
│ ├── gen_wgsl_headers.py # Generate embedded WGSL headers
196+
│ └── test_webgpu_native_ci.sh # CI entry point for native tests
136197
└── test/
137198
├── conftest.py
138-
├── tester.py # Partitioner stages + supported-op list
139-
├── test_build_webgpu.sh # End-to-end build + test
140-
├── test_webgpu_native.cpp # C++ native test runner
141-
├── test_wgsl_codegen.py # Shader codegen check
142-
├── native/ # C++ operator tests
143-
└── ops/ # Python op test suites (flat: test_<op>.py)
144-
├── test_add.py
145-
├── test_rms_norm.py
146-
└── ... # one test_<op>.py per op
199+
├── tester.py # Partitioner and supported-op list
200+
├── test_build_webgpu.sh # End-to-end build and test
201+
├── test_webgpu_native.cpp # Native C++ test runner
202+
├── test_wgsl_codegen.py # Shader drift check
203+
├── native/ # Native C++ operator tests
204+
├── op_tests/ # Codegen case catalog and generator
205+
└── ops/ # Python operator-test modules
147206
```
148207

149208
## Requirements
150209

151-
- **macOS**: Metal-capable GPU
152-
- **Linux**: Vulkan-capable GPU + drivers
153-
- **Build**: CMake 3.19+, conda environment with ExecuTorch installed
210+
- **macOS:** Metal-capable GPU
211+
- **Linux:** Vulkan-capable GPU and drivers
212+
- **Browser:** A WebGPU-enabled browser; the benchmark harness uses Chrome
213+
Canary
214+
- **Build:** CMake 3.19+ and a Python environment with ExecuTorch installed

0 commit comments

Comments
 (0)