RFC: Ratchet V1

[FL33TW00D]

Ratchet as a POC has proved pretty good. A few of the highlights that I am proud of to date:

1. Single model implementation: We thought about datatypes in detail, and as a result we have a single implementation for both full precision and quantised models.
2. Transparent GGUF support: We can pull GGUF models off the hub, parse them on the fly and transcode them into a WebGPU suitable format. This is then cached using IndexedDB. No ONNX conversion, no splitting the model up into 50MB chunks. This ease of use is huge IMO.
3. Hypothesis & Oracle testing: We run 100s of generated tests, and we check our logits against reference implementations for all supported models. This is still a little rough around the edges but is 100% the right direction.
4. Memory Allocator: We've implemented the great memory allocator from Tensorflow.JS and it's fast and efficient with great pooling.

There's a few things we can be less than proud of:

1. Platform support: WebGPU hasn't proliferated as fast and as widely as I would have liked. We need a fast WASM SIMD backend in order to guarantee compatibility, even if it's slower than we would like.
2. Compile time only kernels: This is a big blocker. WGSL has absolutely no preprocessor, which is extremely limiting compared to Metal, which allows you to parameterise of dtypes and all sorts of magic. We need this functionality, because the combinatorics of kernel combinations is going to explode, and we can't be shipping 1e6 pre-generated kernels to users. @philpax had a good idea here of only compiling the kernels the user uses in the program so we don't need to fully move to runtime kernel gen. Unsure what this would look like.
3. Lack of Kernel Fusion: With our current architecture, it's just not possible to fuse even unary ops into their predecessors. This limits our perf.
4. Python Integration: Our oracle tests are a little janky, with python strings embedded in the rust code. I like having the Python code extremely close to the test, but it's pretty bug prone in it's current incarnation.
5. Performance: Our performance is OK, but still loads of room to move. We want to be able to run on shitty intel iGPUs from 2 generations ago, we aren't there yet.

We need an architecture to overcome these limitations. We need to build an abstraction on top of WGSL that provides us the tooling that Apple has with Metal. Easy parameterisation over dtypes, workgroup sizes and more is essential. Furthermore, we need to be able to apply as many sensible fusions as possible.

A WebGPU kernel could be represented in Rust as follows:

pub struct WebGPUKernel {
	bindings: Vec<Binding>, //inputs + outputs + uniform
	globals: Vec<?>, //global functions and shared memory defs
	metadata: Metadata, //required metadata structure
	body: ?, //
}

Given that we are specialised to WebGPU, we should leverage the fact that naga already parses kernels into a nice AST. Once we have this AST, should be able to handle the following case at a minimum:

let gemm = WebGPUKernel { ... }
let gelu = WebGPUKernel { ... }

We intelligently merge bindings, combine globals, and explore the AST for the final store in the GEMM kernel, and modify the AST to wrap the final store with the GELU call.

These are just the ideas that i've had kicking around. Would love input from the kind folks who have been giving me pointers along the way.
@sigma-andex
@ivarflakstad
@philpax
@AmineDiro

[sigma-andex]

My two cents:

1. I agree that having a WASM SIMD backend will be great. However I don't think this needs to be of a high priority right now. Personally I think it is perfectly fine to focus on WebGPU for the moment because:

• It's an early adopter framework
• WebGPU is the future (with more support coming this and probably next year)
• @felladrin has a nice demo of minisearch where he uses different models/frameworks depending on the device/browser. As a user of ratchet for me this would be a perfectly acceptable architecture, i.e. falling back to some other framework in case of no-webgpu.

Again, this is just about priorities/focus. I'm not against it, but I'd rather have a kick-ass WebGPU framework with no cpu support than an ok-ish webgpu framework with cpu support. Also I think that things like 4-bit gguf and performance improvements are more important.

2,3,5) 👍

4. I agree this is janky and at the same time brilliant to use torch as the ground truth here. Also unsure how we could improve that..

[AmineDiro]

Thanks @FL33TW00D for the huge effort! I think that having a full Webgpu ML framework is awesome 🚀. IMO, the GGUF support alone is amazing. This is one of those quality-of-life improvement that will make people choose ratchet for sure.

My one cent on the improvements:
3. Python Integration: Another solution would be to create a Python wrapper library and use python testing instead 🤔 I have seen this pattern in a lot of C++ libs

[FL33TW00D]

"Python wrapper library" - can you expand on this?

For reference, this is what an oracle test looks like for us at the moment: https://github.com/huggingface/ratchet/blob/master/crates/ratchet-models/src/whisper/decoder.rs#L243

[AmineDiro]

"Python wrapper library" - can you expand on this?

The idea is just to use PyO3 to create a python wrapper for ratchet. Once we have that we can use python to test against Pytorch. This would solve the 'jankiness', we would have LSP completion, clean tests,etc... For instance, PyO3 support a #[pymethods] macro to auto generate all struct method as python methods.

The small hiccup is using Python test and not cargo for testing...

[ivarflakstad]

Curveball!
Use tch-rs for ground truth - removing the need for python at all.

[AmineDiro]

@ivarflakstad nice I didn't think about that. I don't know how easy it is to express Pytorch in Rust and if there are some missing missing in this wrapper 🤔

[ivarflakstad]

If you look at the build.rs here you can see that it autogenerates bindings to libtorch (the underlying C++ that pytorch is made up of).

There may be functionality that is not covered by the safe rust wrapper code, but I would be very surprised if it was not enough for our purposes.

[felladrin]

I like what Ratchet achieved so far! The way the GGUF files are loaded is amazing!
I'm eager to use it, so count on me to try the text generation when its NPM package is released.

[felladrin]

I have good news, @sigma-andex and @FL33TW00D!

Ratchet is up and running on MiniSearch!

[FL33TW00D]

Super cool @felladrin nice job!

[FL33TW00D]

Fusion rules for TVM:

• Multiple injective operators can be fused into another injective operator to produce an fused injective operator.
• A reduction operator can be fused with input injective operator to yield a fused injective operator.
• Complex-out-fusable operator can be fused with element-wise operators to its output.

Fusion rules for JITTOR

• Reindex operators may not be fused with preceding meta-operators, as such fusion usually causes performance degradation.
• Reindex-reduce operators may not be fused with following meta-operators.
• Fusion should not create directed cycles between sub-graphs.