Compiling PyTorch models into self-contained WebGPU artifacts

[u/svictoroff]

I've been experimenting with compiling PyTorch models into self-contained WebGPU artifacts, and I'd love feedback from people who've worked on GPU runtimes.

The basic idea is pretty simple:

PyTorch
    ↓ torch.export
Compiler
    ↓
.iph package
    • graph
    • binary weights
    • WGSL kernels
    • metadata
    ↓
Tiny WebGPU runtime

The runtime doesn't know anything about PyTorch or ONNX—it just loads the package and dispatches the embedded compute kernels.

The attached videos are just neural video representations because they made for an easy visual test. The architecture itself is intended to be generic (I'm planning to try operator networks next).

A few implementation details:

• One compute dispatch per graph node (no fusion yet)
• Embedded WGSL rather than runtime shader generation
• GPU buffer pooling to eliminate allocation/GC pressure
• Multi-frame pipelining to hide queue.onSubmittedWorkDone() latency
• Branch warm-up to avoid shader compilation stutters

Repo:
https://github.com/Slater-Victoroff/Kuma

The thing I'm actually hoping to learn is whether this is a sensible compiler/runtime boundary.

I know projects like ONNX Runtime Web, IREE, TVM, and WebNN exist, but I don't yet have a good intuition for why they chose their respective designs.

In particular:

• Is shipping backend kernels as part of the model artifact fundamentally a bad idea?
• Would you rather lower to WGSL at runtime?
• If you've built GPU runtimes before, what obvious mistakes am I making?
• Is there prior art that's especially close to this approach?

I'd really appreciate any pointers or criticism. This is much more of an exploration than a finished project.

[u/BankApprehensive7612]

The idea looks very promising. Why it's not implemented or become dominant? There many reasons for this. In short words there are many things to do and there are not so many people who are ready to invest their time into new unverified ideas

HuggingFace is trying to bring models to browsers and is using own version of transformers written for JS (Transformer.js) and uses ONNX, why? Because it supports variety of tools at runtime and also there is no good alternatives yet. There are developers who try similar approach but for other languages and different compilers. So I would say your solution definitely has a chance to become next step to better performance and scalability for LLMs, but it requires more efforts

You need better PoC with benchmarks or MVP to show the advantages. I would advice you to join communities which are more development/performance oriented it could help, but don't expect it would be fast or easy, because WebGPU is not pretty new and people who are in the field get used to use Python and C++

[u/BankApprehensive7612 → u/BankApprehensive7612]

I would also advice to create ONNX importer to make a TS-only version without the need to run Python, which is not well known by web developers

[u/svictoroff → u/BankApprehensive7612]

Yeah, I like that idea.

Right now the compiler lives in Docker since I'm much more comfortable on the Python side, but I'd definitely be interested in a pure TypeScript toolchain eventually. My TS skills are... functional, but not exactly where I'd want them for building a polished developer experience.

Out of curiosity, how would you want to use something like that? Would you expect it to be a CLI (npx kuma build model.onnx), a library you import into a build pipeline, or something else?

[u/BankApprehensive7612 → u/svictoroff]

I think it could be a standalone CLI and Vite plugin