Make your ZeroGPU Spaces go brrr with ahead-of-time compilation

Nvidia H200 hardware in Hugging Face Spaces without keeping a GPU locked for idle traffic. It’s efficient, flexible, and ideal for demos but it doesn’t always make full use of everything the GPU and CUDA stack can offer. Generating images or videos can take a significant amount of time. Being able to squeeze out more performance, taking advantage of the H200 hardware, does matter in this case.

This is where PyTorch ahead-of-time (AoT) compilation comes in. Instead of compiling models on the fly (which doesn’t play nicely with ZeroGPU’s short-lived processes), AoT lets you optimize once and reload instantly.

The result: snappier demos and a smoother experience, with speedups ranging from 1.3×–1.8× on models like Flux, Wan, and LTX 🔥

In this post, we’ll show how to wire up Ahead-of-Time (AoT) compilation in ZeroGPU Spaces. We'll explore advanced tricks like FP8 quantization and dynamic shapes, and share working demos you can try right away.

What is ZeroGPU

Spaces is a platform powered by Hugging Face that allows ML practitioners to easily publish demo apps. Typical demo apps on Spaces end up reserving a GPU for the Space during its entire lifetime – even when it has no user activity.

ZeroGPU takes a just-in-time approach to GPU initialization. Instead of setting up the main process on CUDA, it automatically forks the process, sets it up on CUDA, runs the GPU tasks, and finally kills the fork when the GPU needs to be released. This means that when the app does not receive traffic, it doesn't use any GPU, but can use multiple GPUs as needed to perform tasks concurrently.

PyTorch compilation

Modern ML frameworks like PyTorch have the concept of compilation that can be used to optimize model latency. PyTorch (from 2.0 onwards) has two major interfaces:

• Just-in-time with torch.compile
• Ahead-of-time with torch.export + AOTInductor

On ZeroGPU, torch.compile can’t efficiently re-use compilation because the process is freshly spun up for almost every task. This is where ahead-of-time (AoT) compilation shines. With AoT, we can export a compiled model once, save it, and later reload it instantly in any process.

Ahead-of-time compilation on ZeroGPU

Compiling a model ahead-of-time involves multiple steps:

1. Capture inputs: Use spaces.aoti_capture to intercept input arguments.
2. Export: Convert the model to a PyTorch ExportedProgram using torch.export.export.
3. Compile: Use spaces.aoti_compile to generate an AoT-compiled binary.
4. Apply: Bind the compiled component back to the original pipeline using spaces.aoti_apply.

This workflow ensures that the heavy computational components (like the transformer in diffusion models) are fully optimized for the H200 hardware without the overhead of runtime compilation.