torch-webgpu

PyTorch compiler and WebGPU runtime, capable of running LLMs like LLama 3.2 3B

Installation

pip install torch-webgpu

Documentation

https://torch-webgpu.maczan.pl

Supported platforms

• Linux (x86_64)
• macOS (Apple Silicon)
• Windows (x86_64)

Use

In Python:

from torch_webgpu import webgpu_backend

And now you can use @torch.compile(backend=webgpu_backend), device="webgpu", to="webgpu" to run and compile PyTorch on a real WebGPU!

FAQ

Why?

WebGPU promises to run everywhere - on almost every hardware - and becomes well supported in web browser. This project is a bridge between PyTorch world and WebGPU world

There is "web" in "WebGPU", so does it mean that I can run PyTorch in a browser now?

This is a step towards running PyTorch in a browser. The next step is to run PyTorch inside a browser. I am actively researching how to do it - if this topic excites you too, contact me on Twitter or open an Issue in this GitHub repo

How serious are you about this project? Is it a research or PoC in mind or are you going to make it production quality?

Once we hit version 1.0.0, torch-webgpu will be a production-ready PyTorch backend. WebGPU is an exciting, emerging technology. As of Nov 2025 all major browsers support WebGPU. I think that it's highly important to build a bridge between PyTorch and WebGPU.

Will you upstream WebGPU backend to PyTorch or keep it out-of-tree forever?

We'll see, ideally I'd see it as a part of PyTorch core, but we need to get a very high quality first to allow ourselves to ask PyTorch maintainers about it

Contributor policy

I have a very little time and need to be picky about contributions, so please make sure you contribute code that is:

• well thought
• covered with unit tests
• you understand everything what you wrote
• as concise as possible - I can't handle too big PRs, sorry!

Use LLM at your discretion, but provide exhaustive explanation of what you built and why. Write it by yourself to show that you really understand

I can understand if that sounds too picky, but since I build this project after hours, I need to cut any additional noise. Sorry and thanks for understanding!

I don't like X about this project

That's ok. The main goal here is to build a bridge (for community) and learn ML compilers in depth (for me). The project moves regularly, at its own pace. Things improve, cover more use cases, get more tests, get rethinked and rewrote. A journey, insights and learning over a raw development velocity. That's a tradeoff I choose

I wish you moved faster

You can fund the project to give me more spare time to work on it. My email: github@maczan.pl

Did AI built it?

The project started 26 Oct 2025. I have been coding it by hand and learning a lot about PyTorch internals and ML compilation in general. Once I made the project to the point where you could compile and run MLP on WebGPU, on 10 Jan 2026 I started to generate many missing ops using AI agents. In just 2 days, AI boosted the project from compiling and running MLPs to compiling and running LLMs ❤️

Open a GitHub issue if you have more questions. Thanks and let's build this bridge!

Ops support

Many of important ops are implemented. If any is missing, feel free to open a PR or an issue. Thanks!

Device / to

• CPU <-> WebGPU
• CUDA <-> WebGPU
• MPS <-> WebGPU
• Intel Gaudi <-> WebGPU
• XLA <-> WebGPU

TODOs

• performance wasn't a priority yet
• only float32 supported
• wgpu::Queue.Submit() handled synchronously
• some ops fallback to CPU
• add more compiler optimizations
• get high performance without platform specific (CUDA, MPS, ROCm) kernels. Five ingredients should be enough to get there - PyTorch, Python, C++, WGSL shaders and WebGPU runtime. Currently, torch-webpgu uses Google Dawn
• implement missing ops

Resources

• Ascend's NPU backend for PyTorch https://github.com/ascend/pytorch
• Elie's WebGPU guide https://eliemichel.github.io/LearnWebGPU/index.html
• WGSL spec https://www.w3.org/TR/WGSL/
• PyTorch PrivateUse1 custom backend docs as a reference https://docs.pytorch.org/tutorials/advanced/privateuseone.html https://docs.pytorch.org/tutorials/advanced/extend_dispatcher.html https://docs.pytorch.org/tutorials/advanced/dispatcher
• https://www.nuss-and-bolts.com/p/optimizing-a-webgpu-matmul-kernel
• https://webgpufundamentals.org/webgpu/lessons/webgpu-compute-shaders.html

Note: This project is unrelated to webgpu-torch, which is a neat PyTorch reimplementation in TypeScript targeting WebGPU

Dev resources

Build from source (only for development)

1. Clone this repo
2. Build Dawn: ./scripts/build-dawn.sh (or set DAWN_PREFIX to your Dawn installation)
3. Build: ./build.sh

C++ unit tests

0. Remember to rebuild your code before testing - ./build.sh
1. chmod +x build-ctests.sh run-ctests.sh
2. Update build-ctests.sh with your paths
3. rm -rf build/ctests
4. ./build-ctests.sh
5. ./run-ctests.sh

C++ benchmarks

0. Remember to rebuild your code before testing - ./build.sh and optionally log in to your wandb.ai account
1. chmod +x build-benchmark.sh run-benchmark.sh
2. Update build-benchmark.sh with your paths
3. rm -rf build/benchmarks
4. ./build-benchmark.sh
5. ./run-benchmark.sh

Python unit tests

0. Remember to rebuild your code before testing - ./build.sh
1. pytest tests to run all tests. pytest tests/ops/test_cos.py to run a chosen test file, like here we test cosinus

Cite

If you use this software, please cite it as below.

@software{Maczan_torch-webgpu_2025,
author = {Maczan, Jędrzej Paweł},
month = oct,
title = {{torch-webgpu - PyTorch compiler and WebGPU runtime}},
url = {https://github.com/jmaczan/torch-webgpu},
version = {1.0.0},
year = {2025}
}

Credits

Jędrzej Maczan, 2025 - ∞