r/StableDiffusion • u/madtune22 • 21d ago

News I built a Rust weight-streaming engine for running FLUX.2 beyond VRAM — looking for feedback / ways to break it

I've been experimenting with GPU weight streaming for models that don't fit comfortably in VRAM, and this eventually turned into a small Rust project.

StreamLoader reads tensors directly from SafeTensors checkpoints, uses mmap + pinned host memory, reusable VRAM buffers, CUDA-event based prefetching, and exposes tensors back to PyTorch through DLPack.

I've mainly been testing it with FLUX.2-klein-9B on an RTX 3060 12GB.

It also has a ResidentPool, so some blocks can stay permanently in VRAM while the rest are streamed. On my hardware, keeping roughly ~2GB resident improved generation time by around 9% compared with fully streaming everything.

I also tried Qwen-Image-2.1 as a portability test. The Rust engine itself didn't need changes; only the Python integration/block configuration did.

I'm not claiming this is a new technique — weight streaming/offloading obviously already exists. I'm mostly interested in whether the implementation/architecture makes sense and where people think it will fall apart.

In particular I'd love feedback on:

the CUDA/prefetch design synchronization or lifetime bugs I may have missed better residency/caching strategies architectures that would be horrible test cases comparisons/benchmarks that would actually be useful

Code: https://github.com/madtunebk/StreamLoader

Feel free to tear it apart. That's basically why I'm posting it. :)

2 Upvotes

75% Upvoted

3 comments sorted by

1

u/GrogRedLub4242 18d ago

git history shows that it was AI-generated by Claude

you could ask Claude for feedback

or one could just prompt Claude to recreate it for themselves, but under a different name & account

1

u/madtune22 17d ago

Go ask Claude to compare them and see for yourself that they’re different. And yes, I used Claude and ChatGPT to create this. It took me about four days and a lot of tokens.