{ "@context": "https://schema.org", "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2025-09-02-make-your-zerogpu-spaces-go-brrr-with-ahead-of-time/#webpage", "@type": [ "WebPage", "ArchiveComponent" ], "articleBody": "ZeroGPU lets anyone spin up powerful **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.\n\nThis 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.\n\n**The result**: snappier demos and a smoother experience, with speedups ranging from **1.3×–1.8×** on models like Flux, Wan, and LTX 🔥\n\nIn 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. If you cannot wait, we invite you to check out some ZeroGPU-powered demos on the [zerogpu-aoti](https://huggingface.co/zerogpu-aoti) organization.\n\n> [Pro](https://huggingface.co/pro) users and [Team / Enterprise](https://huggingface.co/enterprise) org members can create ZeroGPU Spaces, while anyone can freely use them (Pro, Team and Enterprise users get **8x** more ZeroGPU quota)\n\n## Table of Contents\n\n- [What is ZeroGPU](https://huggingface.co/blog/zerogpu-aoti#what-is-zerogpu)\n- [PyTorch compilation](https://huggingface.co/blog/zerogpu-aoti#pytorch-compilation)\n- [Ahead-of-time compilation on ZeroGPU](https://huggingface.co/blog/zerogpu-aoti#ahead-of-time-compilation-on-zerogpu)\n- [Gotchas](https://huggingface.co/blog/zerogpu-aoti#gotchas)\n - [Quantization](https://huggingface.co/blog/zerogpu-aoti#quantization)\n - [Dynamic shapes](https://huggingface.co/blog/zerogpu-aoti#dynamic-shapes)\n - [Multi-compile / shared weights](https://huggingface.co/blog/zerogpu-aoti#multi-compile--shared-weights)\n - [FlashAttention-3](https://huggingface.co/blog/zerogpu-aoti#flashattention-3)\n - [Regional compilation](https://huggingface.co/blog/zerogpu-aoti#regional-compilation)\n - [Use a compiled graph from the Hub](https://huggingface.co/blog/zerogpu-aoti#use-a-compiled-graph-from-the-hub)\n- [AoT compiled ZeroGPU Spaces demos](https://huggingface.co/blog/zerogpu-aoti#aot-compiled-zerogpu-spaces-demos)\n- [Conclusion](https://huggingface.co/blog/zerogpu-aoti#conclusion)\n- [Resources](https://huggingface.co/blog/zerogpu-aoti#resources)\n\n## What is ZeroGPU\n\n[Spaces](https://huggingface.co/spaces) is a platform powered by Hugging Face that allows ML practitioners to easily publish demo apps.\n\nTypical demo apps on Spaces look like:\n\n```python\nimport gradio as gr\nfrom diffusers import DiffusionPipeline\n\npipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n\ndef generate(prompt):\n return pipe(prompt).images\n\ngr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\nThis works great, but ends up reserving a GPU for the Space during its entire lifetime – even when it has no user activity.\n\nWhen executing `.to('cuda')` on this line:\n\n```python\npipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n```\n\nPyTorch initializes the NVIDIA driver, which sets up the process on CUDA forever. This is not very resource-efficient given that app traffic is not perfectly smooth, but is rather extremely sparse and spiky.\n\nZeroGPU 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.\n\nThis means that:\n\n- When the app does not receive traffic, it doesn't use any GPU\n- When it is actually performing a task, it will use one GPU\n- It can use multiple GPUs as needed to perform tasks concurrently\n\nThanks to the Python `spaces` package, the only code change needed to get this behaviour is as follows:\n\n```diff\n import gradio as gr\n+ import spaces\n from diffusers import DiffusionPipeline\n\n pipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n\n+ @spaces.GPU\n def generate(prompt):\n return pipe(prompt).images\n\n gr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\nBy importing `spaces` and adding the `@spaces.GPU` decorator, we:\n\n- Intercept PyTorch API calls to postpone CUDA operations\n- Make the decorated function run in a fork when later called\n- (Call an internal API to make the right device visible to the fork but this is not in the scope of this blogpost)\n\n> ZeroGPU currently allocates an [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/#h200-mig-profiles) slice of H200 (`3g.71gb` profile). Additional MIG sizes including full slice (`7g.141gb` profile) will come in late 2025.\n\n## PyTorch compilation\n\nModern ML frameworks like PyTorch and JAX have the concept of *compilation* that can be used to optimize model latency or inference time. Behind the scenes, compilation applies a series of (often hardware-dependent) optimization steps such as operator fusion, constant folding, etc.\n\nPyTorch (from 2.0 onwards) currently has two major interfaces for compilation:\n\n- Just-in-time with `torch.compile`\n- Ahead-of-time with `torch.export` + `AOTInductor`\n\n[`torch.compile`](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) works great in standard environments: it compiles your model the first time it runs, and reuses the optimized version for subsequent calls.\n\nHowever, on ZeroGPU, given that the process is freshly spun up for (almost) every GPU task, it means that `torch.compile` can’t efficiently re-use compilation and is thus forced to rely on its [filesystem cache](https://docs.pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html#modular-caching-of-torchdynamo-torchinductor-and-triton) to restore compiled models. Depending on the model being compiled, this process takes from a few dozen seconds to a couple of minutes, which is way too much for practical GPU tasks in Spaces.\n\nThis is where **ahead-of-time (AoT) compilation** shines.\n\nWith AoT, we can export a compiled model once, save it, and later reload it instantly in any process, which is exactly what we need for ZeroGPU. This helps us reduce framework overhead and also eliminates cold-start timings typically incurred in just-in-time compilation.\n\nBut how can we do ahead-of-time compilation on ZeroGPU? Let’s dive in.\n\n## Ahead-of-time compilation on ZeroGPU\n\nLet's go back to our ZeroGPU base example and unpack what we need to enable AoT compilation. For the purpose of this demo, we will use the `black-forest-labs/FLUX.1-dev` model:\n\n```python\nimport gradio as gr\nimport spaces\nimport torch\nfrom diffusers import DiffusionPipeline\n\nMODEL_ID = 'black-forest-labs/FLUX.1-dev'\n\npipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)\npipe.to('cuda')\n\n@spaces.GPU\ndef generate(prompt):\n return pipe(prompt).images\n\ngr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\n> In the discussion below, we only compile the `transformer` component of `pipe` since, in these generative models, the transformer (or more generally, the denoiser) is the most computationally heavy component.\n\nCompiling a model ahead-of-time with PyTorch involves multiple steps:\n\n### 1. Getting example inputs\n\nRecall that we’re going to compile the model *ahead* of time. Therefore, we need to derive example inputs for the model. Note that these are the same kinds of inputs we expect to see during the actual runs. To capture those inputs, we will leverage the `spaces.aoti_capture` helper from the `spaces` package:\n\n```python\nwith spaces.aoti_capture(pipe.transformer) as call:\n pipe(\"arbitrary example prompt\")\n```\n\nWhen used as a context manager, `aoti_capture` intercepts the call to any callable (`pipe.transformer` in our case), prevents it from executing, captures the input arguments that would have been passed to it, and stores their values in `call.args` and `call.kwargs`.\n\n### 2. Exporting the model\n\nNow that we have example args and kwargs for our transformer component, we can export it to a PyTorch [`ExportedProgram`](https://docs.pytorch.org/docs/stable/export.html#torch.export.ExportedProgram) by using [`torch.export.export`](https://docs.pytorch.org/docs/stable/export.html#torch.export.export) utility:\n\n```python\nexported_transformer = torch.export.export(\n pipe.transformer,\n args=call.args,\n kwargs=call.kwargs,\n)\n```\n\nAn exported PyTorch program is a computation graph that represents the tensor computations along with the original model parameter values.\n\n### 3. Compiling the exported model\n\nOnce the model is exported, compiling it is pretty straightforward.\n\nA traditional AoT compilation in PyTorch often requires saving the model on disk so it can be later reloaded. In our case, we’ll leverage a helper function part of the `spaces` package: `spaces.aoti_compile`. It's a tiny wrapper around `torch._inductor.aot_compile` that manages saving and lazy-loading the model as needed. It's meant to be used like this:\n\n```python\ncompiled_transformer = spaces.aoti_compile(exported_transformer)\n```\n\nThis `compiled_transformer` is now an AoT-compiled binary ready to be used for inference.\n\n### 4. Using the compiled model in the pipeline\n\nNow we need to bind our compiled transformer to our original pipeline, i.e., the `pipeline`.\n\nA naive and almost working approach is to simply patch our pipeline like `pipe.transformer = compiled_transformer`. Unfortunately, this approach does not work because it deletes important attributes like `dtype`, `config`, etc. Only patching the `forward` method does not work well either because we are then keeping original model parameters in memory, often leading to OOM errors at runtime.\n\n`spaces` package provides a utility for this, too -- `spaces.aoti_apply`:\n\n```python\nspaces.aoti_apply(compiled_transformer, pipe.transformer)\n```\n\nEt voilà! It will take care of patching `pipe.transformer.forward` with our compiled model, as well as [cleaning old model parameters out of memory](https://pypi-browser.org/package/spaces/spaces-0.40.1-py3-none-any.whl/spaces/zero/torch/aoti.py#L87).\n\n### 5. Wrapping it all together\n\nTo perform the first three steps (intercepting input examples, exporting the model, and compiling it with PyTorch inductor), we need a real GPU. CUDA emulation that you get outside of `@spaces.GPU` function is not enough because compilation is truly hardware-dependent, for instance, relying on micro-benchmark runs to tune the generated code. This is why we need to wrap it all inside a `@spaces.GPU` function and then get our compiled model back to the root of our app. Starting from our original demo code, this gives:\n\n```diff\n import gradio as gr\n import spaces\n import torch\n from diffusers import DiffusionPipeline\n \n MODEL_ID = 'black-forest-labs/FLUX.1-dev'\n \n pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)\n pipe.to('cuda')\n \n+ @spaces.GPU(duration=1500) # maximum duration allowed during startup\n+ def compile_transformer():\n+ with spaces.aoti_capture(pipe.transformer) as call:\n+ pipe(\"arbitrary example prompt\")\n+ \n+ exported = torch.export.export(\n+ pipe.transformer,\n+ args=call.args,\n+ kwargs=call.kwargs,\n+ )\n+ return spaces.aoti_compile(exported)\n+ \n+ compiled_transformer = compile_transformer()\n+ spaces.aoti_apply(compiled_transformer, pipe.transformer)\n \n @spaces.GPU\n def generate(prompt):\n return pipe(prompt).images\n \n gr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\nWith just a dozen lines of additional code, we’ve successfully made our demo quite faster (**1.7x** faster in the case of FLUX.1-dev).\n\nIf you want to learn more about AoT compilation, you can read PyTorch's [AOTInductor tutorial](https://docs.pytorch.org/tutorials/recipes/torch_export_aoti_python.html)\n\n## Gotchas\n\nNow that we have demonstrated the speedups one can realize under the constraints of operating with ZeroGPUs, we will discuss a few gotchas that came up while working with this setup.\n\n### Quantization\n\nAoT can be combined with quantization to deliver even greater speedups. For image and video generation, the FP8 post-training dynamic quantization schemes deliver good speed-quality trade-offs. However, FP8 requires a CUDA compute capability of at least 9.0 to work. Thankfully, for ZeroGPUs, since they’re based on H200s, we can already take advantage of the FP8 quantization schemes.\n\nTo enable FP8 quantization within our AoT compilation workflow, we can leverage the APIs provided by [`torchao`](https://github.com/pytorch/ao) like so:\n\n```diff\n+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig\n\n+ # Quantize the transformer just before the export step.\n+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())\n\nexported_transformer = torch.export.export(\n pipe.transformer,\n args=call.args,\n kwargs=call.kwargs,\n)\n```\n\n(You can find more details about TorchAO [here](https://docs.pytorch.org/ao/stable/index.html).)\n\nAnd we can then proceed with the rest of the steps as outlined above. Using quantization provides another **1.2x** of speedup.\n\n### Dynamic shapes\n\nImages and videos can come in different shapes and sizes. Hence, it’s important to also account for shape dynamism when performing AoT compilation. The primitives provided by `torch.export.export` make it easily configurable to provide what inputs should be treated accordingly for dynamic shapes, as shown below.\n\nFor the case of Flux.1-Dev transformer, changes in different image resolutions will affect two of its `forward` arguments:\n\n- `hidden_states`: The noisy input latents, which the transformer is supposed to denoise. It’s a 3D tensor, representing `batch_size, flattened_latent_dim, embed_dim`. When the batch size is fixed, it’s the `flattened_latent_dim` that will change for any changes made to image resolutions.\n- `img_ids`: A 2D array of encoded pixel coordinates having a shape of `height * width, 3`. In this case, we want to make `height * width` dynamic.\n\nWe start by defining a range in which we want to let the (latent) image resolutions vary. To derive these value ranges, we inspected the shapes of [`hidden_states`](https://github.com/huggingface/diffusers/blob/0ff1aa910cf3d87193af79ec1ae4487be542e872/src/diffusers/pipelines/flux/pipeline_flux.py#L920) in the pipeline with respect to varied image resolutions. The exact values are model-dependent and require manual inspection and some intuition. For Flux.1-Dev, we ended up with:\n\n```python\ntransformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)\n```\n\nWe then define a map of argument names and which dimensions in their input values we expect to be dynamic:\n\n```python\ntransformer_dynamic_shapes = {\n \"hidden_states\": {1: transformer_hidden_dim}, \n \"img_ids\": {0: transformer_hidden_dim},\n}\n```\n\nThen we need to make our dynamic shapes object replicate the structure of our example inputs. The inputs that do not need dynamic shapes must be set to `None`. This can be done very easily with PyTorch [tree\\_map](https://github.com/pytorch/pytorch/blob/2f0de0ff9361ca4f2b1e6f9edbc600b5fb6abcd6/torch/utils/_pytree.py#L1341-L1373) utility:\n\n```python\nfrom torch.utils._pytree import tree_map\n\ndynamic_shapes = tree_map(lambda v: None, call.kwargs)\ndynamic_shapes |= transformer_dynamic_shapes\n```\n\nNow, when performing the export step, we simply supply `transformer_dynamic_shapes` to `torch.export.export`:\n\n```python\nexported_transformer = torch.export.export(\n pipe.transformer,\n args=call.args,\n kwargs=call.kwargs,\n dynamic_shapes=dynamic_shapes,\n)\n```\n\n> Check out [this Space](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic) that shows how to use both quantization and dynamic shapes during the export step.\n\n### Multi-compile / shared weights\n\nDynamic shapes is sometimes not enough when dynamism is too important.\n\nThis is, for instance, the case with the Wan family of video generation models if you want your compiled model to generate different resolutions. One thing can be done in this case: compile one model per resolution while keeping the model parameters shared and dispatching the right one at runtime\n\nHere is a minimal example of this approach: [zerogpu-aoti-multi.py](https://gist.github.com/cbensimon/8dc0ffcd7ee024d91333f6df01907916). You can also see a fully working implementation of this paradigm in the [Wan 2.2 Space](https://huggingface.co/spaces/zerogpu-aoti/wan2-2-fp8da-aoti-faster/blob/main/optimization.py).\n\n### FlashAttention-3\n\nSince the ZeroGPU hardware and CUDA drivers are perfectly compatible with Flash-Attention 3 (FA3), we can use it in our ZeroGPU Spaces to speed things up even further. FA3 works with ahead-of-time compilation. So, this is ideal for our case.\n\nCompiling and building FA3 from source can take several minutes, and this process is hardware-dependent. As users, we wouldn’t want to lose precious ZeroGPU compute hours. This is where Hugging Face [`kernels` library](https://github.com/huggingface/kernels) comes to the rescue. It provides access to pre-built kernels that are compatible for a given hardware. For example, when we try to run:\n\n```python\nfrom kernels import get_kernel\n\nvllm_flash_attn3 = get_kernel(\"kernels-community/vllm-flash-attn3\")\n```\n\nIt tries to load a kernel from the [`kernels-community/vllm-flash-attn3`](https://huggingface.co/kernels-community/vllm-flash-attn3) repository, which is compatible with the current setup. Otherwise, it will error out due to incompatibility issues. Luckily for us, this works seamlessly on the ZeroGPU Spaces. This means we can leverage the power of FA3 on ZeroGPU, using the `kernels` library.\n\nHere is a [fully working example of an FA3 attention processor](https://gist.github.com/sayakpaul/ff715f979793d4d44beb68e5e08ee067#file-fa3_qwen-py) for the Qwen-Image model.\n\n### Regional compilation\n\nSo far, we have been compiling the full model. Depending on the model, full model compilation can lead to significantly long cold start times. Long cold start times make the development experience unpleasant.\n\nWe can also choose to compile *regions* within a model, significantly reducing the cold start times, while retaining almost all the benefits of full model compilation. Regional compilation becomes promising when a model has repeated blocks of computation. A standard language model, for example, has a number of identically structured Transformer blocks.\n\nIn our example, we can compile the repeated blocks of the Flux transformer ahead of time, and propagate the compiled graph to the remaining repeated blocks. The [Flux Transformer](https://github.com/huggingface/diffusers/blob/c2e5ece08bf22d249c62e964f91bc326cf9e3759/src/diffusers/models/transformers/transformer_flux.py) has two kinds of repeated blocks: `FluxTransformerBlock` and `FluxSingleTransformerBlock`.\n\nYou can check out [this Space](https://huggingface.co/spaces/cbensimon/FLUX.1-dev-fa3-aoti/tree/main) for a complete example.\n\n> 💡 For Flux.1-Dev, switching to regional compilation reduces the compilation time from *6 minutes* to just *30 seconds* while delivering identical speedups.\n\n### Use a compiled graph from the Hub\n\nOnce a model (or even a model block) is compiled ahead of time, we can serialize the compiled graph module as an artifact and reuse later. In the context of a ZeroGPU-powered demo on Spaces, this will significantly cut down the demo startup time by skipping the compilation time.\n\nTo keep the storage light, we can just save the compiled model graph without including any model parameters inside the artifact.\n\nCheck out [this collection](https://huggingface.co/collections/zerogpu-aoti/using-compiled-graph-from-the-hub-68c2afcc03de7609f9f91e35) that shows a full workflow of obtaining compiled model graph, pushing it to the Hub, and then using it to build a demo.\n\n## AoT compiled ZeroGPU Spaces demos\n\n### Speedup comparison\n\n- [FLUX.1-dev without AoTI](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-dev-base)\n- [FLUX.1-dev with AoTI and FA3](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-dev-fa3-aoti) (**1.75x** speedup)\n\n### Featured AoTI Spaces\n\n- [FLUX.1 Kontext](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev)\n- [QwenImage Edit](https://huggingface.co/spaces/multimodalart/Qwen-Image-Edit-Fast)\n- [Wan 2.2](https://huggingface.co/spaces/zerogpu-aoti/wan2-2-fp8da-aoti-faster)\n\n### Regional compilation\n\n- [Regional compilation recipe](https://docs.pytorch.org/tutorials/recipes/regional_compilation.html)\n- [Regional compilation with AOT](https://docs.pytorch.org/tutorials/recipes/regional_aot.html)\n- [Native integration in Diffusers](https://huggingface.co/docs/diffusers/main/en/optimization/fp16)\n- [More performance numbers](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/)\n\n## Conclusion\n\nZeroGPU within Hugging Face Spaces is a powerful feature that enables AI builders by providing access to powerful compute. In this post, we showed how users can benefit from PyTorch’s ahead-of-time compilation techniques to speed up their applications that leverage ZeroGPU.\n\nWe demonstrate speedups with Flux.1-Dev, but these techniques are not limited to just this model. Therefore, we encourage you to give these techniques a try and provide us with feedback in this [community discussion](https://huggingface.co/spaces/zerogpu-aoti/README/discussions/1).\n\n## Resources\n\n- Visit our [ZeroGPU-AOTI org on the Hub](https://huggingface.co/zerogpu-aoti) to refer to a collection of demos that leverage the techniques discussed in this post.\n- Browse `spaces.aoti_*` APIs [source code](https://pypi-browser.org/package/spaces/spaces-0.40.1-py3-none-any.whl/spaces/zero/torch/aoti.py) to learn more about the interface\n- Check out [Kernels Community org on the hub](https://huggingface.co/kernels-community)\n- Learn more about regional compilation from [here](https://huggingface.co/blog/pytorch.org/tutorials/recipes/regional_compilation.html)\n- Upgrade to [Pro](https://huggingface.co/pro) on Hugging Face to create your own ZeroGPU Spaces (and get 25 minutes of H200 usage every day)\n\n*Acknowledgements: Thanks to ChunTe Lee for creating an awesome thumbnail for this post. Thanks to Pedro and Vaibhav for providing feedback on the post. Thanks to Angela Yi from the PyTorch team for helping us with AOT guidance.*\n", "authors": [], "captured_at": "2026-10-04T12:50:30.804479567+00:00", "category": "ai", "content_html": "

ZeroGPU lets anyone spin up powerful 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.

\n

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.

\n

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

\n

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. If you cannot wait, we invite you to check out some ZeroGPU-powered demos on the zerogpu-aoti organization.

\n
\n

Pro users and Team / Enterprise org members can create ZeroGPU Spaces, while anyone can freely use them (Pro, Team and Enterprise users get 8x more ZeroGPU quota)

\n
\n

Table of Contents

\n\n

What is ZeroGPU

\n

Spaces is a platform powered by Hugging Face that allows ML practitioners to easily publish demo apps.

\n

Typical demo apps on Spaces look like:

\n
import gradio as gr\nfrom diffusers import DiffusionPipeline\n\npipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n\ndef generate(prompt):\n    return pipe(prompt).images\n\ngr.Interface(generate, "text", "gallery").launch()\n
\n

This works great, but ends up reserving a GPU for the Space during its entire lifetime – even when it has no user activity.

\n

When executing .to('cuda') on this line:

\n
pipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n
\n

PyTorch initializes the NVIDIA driver, which sets up the process on CUDA forever. This is not very resource-efficient given that app traffic is not perfectly smooth, but is rather extremely sparse and spiky.

\n

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.

\n

This means that:

\n\n

Thanks to the Python spaces package, the only code change needed to get this behaviour is as follows:

\n
  import gradio as gr\n+ import spaces\n  from diffusers import DiffusionPipeline\n\n  pipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n\n+ @spaces.GPU\n  def generate(prompt):\n      return pipe(prompt).images\n\n  gr.Interface(generate, "text", "gallery").launch()\n
\n

By importing spaces and adding the @spaces.GPU decorator, we:

\n\n
\n

ZeroGPU currently allocates an MIG slice of H200 (3g.71gb profile). Additional MIG sizes including full slice (7g.141gb profile) will come in late 2025.

\n
\n

PyTorch compilation

\n

Modern ML frameworks like PyTorch and JAX have the concept of compilation that can be used to optimize model latency or inference time. Behind the scenes, compilation applies a series of (often hardware-dependent) optimization steps such as operator fusion, constant folding, etc.

\n

PyTorch (from 2.0 onwards) currently has two major interfaces for compilation:

\n\n

torch.compile works great in standard environments: it compiles your model the first time it runs, and reuses the optimized version for subsequent calls.

\n

However, on ZeroGPU, given that the process is freshly spun up for (almost) every GPU task, it means that torch.compile can’t efficiently re-use compilation and is thus forced to rely on its filesystem cache to restore compiled models. Depending on the model being compiled, this process takes from a few dozen seconds to a couple of minutes, which is way too much for practical GPU tasks in Spaces.

\n

This is where ahead-of-time (AoT) compilation shines.

\n

With AoT, we can export a compiled model once, save it, and later reload it instantly in any process, which is exactly what we need for ZeroGPU. This helps us reduce framework overhead and also eliminates cold-start timings typically incurred in just-in-time compilation.

\n

But how can we do ahead-of-time compilation on ZeroGPU? Let’s dive in.

\n

Ahead-of-time compilation on ZeroGPU

\n

Let's go back to our ZeroGPU base example and unpack what we need to enable AoT compilation. For the purpose of this demo, we will use the black-forest-labs/FLUX.1-dev model:

\n
import gradio as gr\nimport spaces\nimport torch\nfrom diffusers import DiffusionPipeline\n\nMODEL_ID = 'black-forest-labs/FLUX.1-dev'\n\npipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)\npipe.to('cuda')\n\n@spaces.GPU\ndef generate(prompt):\n    return pipe(prompt).images\n\ngr.Interface(generate, "text", "gallery").launch()\n
\n
\n

In the discussion below, we only compile the transformer component of pipe since, in these generative models, the transformer (or more generally, the denoiser) is the most computationally heavy component.

\n
\n

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

\n

1. Getting example inputs

\n

Recall that we’re going to compile the model ahead of time. Therefore, we need to derive example inputs for the model. Note that these are the same kinds of inputs we expect to see during the actual runs. To capture those inputs, we will leverage the spaces.aoti_capture helper from the spaces package:

\n
with spaces.aoti_capture(pipe.transformer) as call:\n    pipe("arbitrary example prompt")\n
\n

When used as a context manager, aoti_capture intercepts the call to any callable (pipe.transformer in our case), prevents it from executing, captures the input arguments that would have been passed to it, and stores their values in call.args and call.kwargs.

\n

2. Exporting the model

\n

Now that we have example args and kwargs for our transformer component, we can export it to a PyTorch ExportedProgram by using torch.export.export utility:

\n
exported_transformer = torch.export.export(\n    pipe.transformer,\n    args=call.args,\n    kwargs=call.kwargs,\n)\n
\n

An exported PyTorch program is a computation graph that represents the tensor computations along with the original model parameter values.

\n

3. Compiling the exported model

\n

Once the model is exported, compiling it is pretty straightforward.

\n

A traditional AoT compilation in PyTorch often requires saving the model on disk so it can be later reloaded. In our case, we’ll leverage a helper function part of the spaces package: spaces.aoti_compile. It's a tiny wrapper around torch._inductor.aot_compile that manages saving and lazy-loading the model as needed. It's meant to be used like this:

\n
compiled_transformer = spaces.aoti_compile(exported_transformer)\n
\n

This compiled_transformer is now an AoT-compiled binary ready to be used for inference.

\n

4. Using the compiled model in the pipeline

\n

Now we need to bind our compiled transformer to our original pipeline, i.e., the pipeline.

\n

A naive and almost working approach is to simply patch our pipeline like pipe.transformer = compiled_transformer. Unfortunately, this approach does not work because it deletes important attributes like dtype, config, etc. Only patching the forward method does not work well either because we are then keeping original model parameters in memory, often leading to OOM errors at runtime.

\n

spaces package provides a utility for this, too -- spaces.aoti_apply:

\n
spaces.aoti_apply(compiled_transformer, pipe.transformer)\n
\n

Et voilà! It will take care of patching pipe.transformer.forward with our compiled model, as well as cleaning old model parameters out of memory.

\n

5. Wrapping it all together

\n

To perform the first three steps (intercepting input examples, exporting the model, and compiling it with PyTorch inductor), we need a real GPU. CUDA emulation that you get outside of @spaces.GPU function is not enough because compilation is truly hardware-dependent, for instance, relying on micro-benchmark runs to tune the generated code. This is why we need to wrap it all inside a @spaces.GPU function and then get our compiled model back to the root of our app. Starting from our original demo code, this gives:

\n
  import gradio as gr\n  import spaces\n  import torch\n  from diffusers import DiffusionPipeline\n  \n  MODEL_ID = 'black-forest-labs/FLUX.1-dev'\n  \n  pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)\n  pipe.to('cuda')\n  \n+ @spaces.GPU(duration=1500) # maximum duration allowed during startup\n+ def compile_transformer():\n+     with spaces.aoti_capture(pipe.transformer) as call:\n+         pipe("arbitrary example prompt")\n+ \n+     exported = torch.export.export(\n+         pipe.transformer,\n+         args=call.args,\n+         kwargs=call.kwargs,\n+     )\n+     return spaces.aoti_compile(exported)\n+ \n+ compiled_transformer = compile_transformer()\n+ spaces.aoti_apply(compiled_transformer, pipe.transformer)\n  \n  @spaces.GPU\n  def generate(prompt):\n      return pipe(prompt).images\n  \n  gr.Interface(generate, "text", "gallery").launch()\n
\n

With just a dozen lines of additional code, we’ve successfully made our demo quite faster (1.7x faster in the case of FLUX.1-dev).

\n

If you want to learn more about AoT compilation, you can read PyTorch's AOTInductor tutorial

\n

Gotchas

\n

Now that we have demonstrated the speedups one can realize under the constraints of operating with ZeroGPUs, we will discuss a few gotchas that came up while working with this setup.

\n

Quantization

\n

AoT can be combined with quantization to deliver even greater speedups. For image and video generation, the FP8 post-training dynamic quantization schemes deliver good speed-quality trade-offs. However, FP8 requires a CUDA compute capability of at least 9.0 to work. Thankfully, for ZeroGPUs, since they’re based on H200s, we can already take advantage of the FP8 quantization schemes.

\n

To enable FP8 quantization within our AoT compilation workflow, we can leverage the APIs provided by torchao like so:

\n
+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig\n\n+ # Quantize the transformer just before the export step.\n+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())\n\nexported_transformer = torch.export.export(\n    pipe.transformer,\n    args=call.args,\n    kwargs=call.kwargs,\n)\n
\n

(You can find more details about TorchAO here.)

\n

And we can then proceed with the rest of the steps as outlined above. Using quantization provides another 1.2x of speedup.

\n

Dynamic shapes

\n

Images and videos can come in different shapes and sizes. Hence, it’s important to also account for shape dynamism when performing AoT compilation. The primitives provided by torch.export.export make it easily configurable to provide what inputs should be treated accordingly for dynamic shapes, as shown below.

\n

For the case of Flux.1-Dev transformer, changes in different image resolutions will affect two of its forward arguments:

\n\n

We start by defining a range in which we want to let the (latent) image resolutions vary. To derive these value ranges, we inspected the shapes of hidden_states in the pipeline with respect to varied image resolutions. The exact values are model-dependent and require manual inspection and some intuition. For Flux.1-Dev, we ended up with:

\n
transformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)\n
\n

We then define a map of argument names and which dimensions in their input values we expect to be dynamic:

\n
transformer_dynamic_shapes = {\n    "hidden_states": {1: transformer_hidden_dim}, \n    "img_ids": {0: transformer_hidden_dim},\n}\n
\n

Then we need to make our dynamic shapes object replicate the structure of our example inputs. The inputs that do not need dynamic shapes must be set to None. This can be done very easily with PyTorch tree_map utility:

\n
from torch.utils._pytree import tree_map\n\ndynamic_shapes = tree_map(lambda v: None, call.kwargs)\ndynamic_shapes |= transformer_dynamic_shapes\n
\n

Now, when performing the export step, we simply supply transformer_dynamic_shapes to torch.export.export:

\n
exported_transformer = torch.export.export(\n    pipe.transformer,\n    args=call.args,\n    kwargs=call.kwargs,\n    dynamic_shapes=dynamic_shapes,\n)\n
\n
\n

Check out this Space that shows how to use both quantization and dynamic shapes during the export step.

\n
\n

Multi-compile / shared weights

\n

Dynamic shapes is sometimes not enough when dynamism is too important.

\n

This is, for instance, the case with the Wan family of video generation models if you want your compiled model to generate different resolutions. One thing can be done in this case: compile one model per resolution while keeping the model parameters shared and dispatching the right one at runtime

\n

Here is a minimal example of this approach: zerogpu-aoti-multi.py. You can also see a fully working implementation of this paradigm in the Wan 2.2 Space.

\n

FlashAttention-3

\n

Since the ZeroGPU hardware and CUDA drivers are perfectly compatible with Flash-Attention 3 (FA3), we can use it in our ZeroGPU Spaces to speed things up even further. FA3 works with ahead-of-time compilation. So, this is ideal for our case.

\n

Compiling and building FA3 from source can take several minutes, and this process is hardware-dependent. As users, we wouldn’t want to lose precious ZeroGPU compute hours. This is where Hugging Face kernels library comes to the rescue. It provides access to pre-built kernels that are compatible for a given hardware. For example, when we try to run:

\n
from kernels import get_kernel\n\nvllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")\n
\n

It tries to load a kernel from the kernels-community/vllm-flash-attn3 repository, which is compatible with the current setup. Otherwise, it will error out due to incompatibility issues. Luckily for us, this works seamlessly on the ZeroGPU Spaces. This means we can leverage the power of FA3 on ZeroGPU, using the kernels library.

\n

Here is a fully working example of an FA3 attention processor for the Qwen-Image model.

\n

Regional compilation

\n

So far, we have been compiling the full model. Depending on the model, full model compilation can lead to significantly long cold start times. Long cold start times make the development experience unpleasant.

\n

We can also choose to compile regions within a model, significantly reducing the cold start times, while retaining almost all the benefits of full model compilation. Regional compilation becomes promising when a model has repeated blocks of computation. A standard language model, for example, has a number of identically structured Transformer blocks.

\n

In our example, we can compile the repeated blocks of the Flux transformer ahead of time, and propagate the compiled graph to the remaining repeated blocks. The Flux Transformer has two kinds of repeated blocks: FluxTransformerBlock and FluxSingleTransformerBlock.

\n

You can check out this Space for a complete example.

\n
\n

💡 For Flux.1-Dev, switching to regional compilation reduces the compilation time from 6 minutes to just 30 seconds while delivering identical speedups.

\n
\n

Use a compiled graph from the Hub

\n

Once a model (or even a model block) is compiled ahead of time, we can serialize the compiled graph module as an artifact and reuse later. In the context of a ZeroGPU-powered demo on Spaces, this will significantly cut down the demo startup time by skipping the compilation time.

\n

To keep the storage light, we can just save the compiled model graph without including any model parameters inside the artifact.

\n

Check out this collection that shows a full workflow of obtaining compiled model graph, pushing it to the Hub, and then using it to build a demo.

\n

AoT compiled ZeroGPU Spaces demos

\n

Speedup comparison

\n\n

Featured AoTI Spaces

\n\n

Regional compilation

\n\n

Conclusion

\n

ZeroGPU within Hugging Face Spaces is a powerful feature that enables AI builders by providing access to powerful compute. In this post, we showed how users can benefit from PyTorch’s ahead-of-time compilation techniques to speed up their applications that leverage ZeroGPU.

\n

We demonstrate speedups with Flux.1-Dev, but these techniques are not limited to just this model. Therefore, we encourage you to give these techniques a try and provide us with feedback in this community discussion.

\n

Resources

\n\n

Acknowledgements: Thanks to ChunTe Lee for creating an awesome thumbnail for this post. Thanks to Pedro and Vaibhav for providing feedback on the post. Thanks to Angela Yi from the PyTorch team for helping us with AOT guidance.

\n", "content_markdown": "ZeroGPU lets anyone spin up powerful **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.\n\nThis 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.\n\n**The result**: snappier demos and a smoother experience, with speedups ranging from **1.3×–1.8×** on models like Flux, Wan, and LTX 🔥\n\nIn 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. If you cannot wait, we invite you to check out some ZeroGPU-powered demos on the [zerogpu-aoti](https://huggingface.co/zerogpu-aoti) organization.\n\n> [Pro](https://huggingface.co/pro) users and [Team / Enterprise](https://huggingface.co/enterprise) org members can create ZeroGPU Spaces, while anyone can freely use them (Pro, Team and Enterprise users get **8x** more ZeroGPU quota)\n\n## Table of Contents\n\n- [What is ZeroGPU](https://huggingface.co/blog/zerogpu-aoti#what-is-zerogpu)\n- [PyTorch compilation](https://huggingface.co/blog/zerogpu-aoti#pytorch-compilation)\n- [Ahead-of-time compilation on ZeroGPU](https://huggingface.co/blog/zerogpu-aoti#ahead-of-time-compilation-on-zerogpu)\n- [Gotchas](https://huggingface.co/blog/zerogpu-aoti#gotchas)\n - [Quantization](https://huggingface.co/blog/zerogpu-aoti#quantization)\n - [Dynamic shapes](https://huggingface.co/blog/zerogpu-aoti#dynamic-shapes)\n - [Multi-compile / shared weights](https://huggingface.co/blog/zerogpu-aoti#multi-compile--shared-weights)\n - [FlashAttention-3](https://huggingface.co/blog/zerogpu-aoti#flashattention-3)\n - [Regional compilation](https://huggingface.co/blog/zerogpu-aoti#regional-compilation)\n - [Use a compiled graph from the Hub](https://huggingface.co/blog/zerogpu-aoti#use-a-compiled-graph-from-the-hub)\n- [AoT compiled ZeroGPU Spaces demos](https://huggingface.co/blog/zerogpu-aoti#aot-compiled-zerogpu-spaces-demos)\n- [Conclusion](https://huggingface.co/blog/zerogpu-aoti#conclusion)\n- [Resources](https://huggingface.co/blog/zerogpu-aoti#resources)\n\n## What is ZeroGPU\n\n[Spaces](https://huggingface.co/spaces) is a platform powered by Hugging Face that allows ML practitioners to easily publish demo apps.\n\nTypical demo apps on Spaces look like:\n\n```python\nimport gradio as gr\nfrom diffusers import DiffusionPipeline\n\npipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n\ndef generate(prompt):\n return pipe(prompt).images\n\ngr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\nThis works great, but ends up reserving a GPU for the Space during its entire lifetime – even when it has no user activity.\n\nWhen executing `.to('cuda')` on this line:\n\n```python\npipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n```\n\nPyTorch initializes the NVIDIA driver, which sets up the process on CUDA forever. This is not very resource-efficient given that app traffic is not perfectly smooth, but is rather extremely sparse and spiky.\n\nZeroGPU 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.\n\nThis means that:\n\n- When the app does not receive traffic, it doesn't use any GPU\n- When it is actually performing a task, it will use one GPU\n- It can use multiple GPUs as needed to perform tasks concurrently\n\nThanks to the Python `spaces` package, the only code change needed to get this behaviour is as follows:\n\n```diff\n import gradio as gr\n+ import spaces\n from diffusers import DiffusionPipeline\n\n pipe = DiffusionPipeline.from_pretrained(...).to('cuda')\n\n+ @spaces.GPU\n def generate(prompt):\n return pipe(prompt).images\n\n gr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\nBy importing `spaces` and adding the `@spaces.GPU` decorator, we:\n\n- Intercept PyTorch API calls to postpone CUDA operations\n- Make the decorated function run in a fork when later called\n- (Call an internal API to make the right device visible to the fork but this is not in the scope of this blogpost)\n\n> ZeroGPU currently allocates an [MIG](https://docs.nvidia.com/datacenter/tesla/mig-user-guide/#h200-mig-profiles) slice of H200 (`3g.71gb` profile). Additional MIG sizes including full slice (`7g.141gb` profile) will come in late 2025.\n\n## PyTorch compilation\n\nModern ML frameworks like PyTorch and JAX have the concept of *compilation* that can be used to optimize model latency or inference time. Behind the scenes, compilation applies a series of (often hardware-dependent) optimization steps such as operator fusion, constant folding, etc.\n\nPyTorch (from 2.0 onwards) currently has two major interfaces for compilation:\n\n- Just-in-time with `torch.compile`\n- Ahead-of-time with `torch.export` + `AOTInductor`\n\n[`torch.compile`](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) works great in standard environments: it compiles your model the first time it runs, and reuses the optimized version for subsequent calls.\n\nHowever, on ZeroGPU, given that the process is freshly spun up for (almost) every GPU task, it means that `torch.compile` can’t efficiently re-use compilation and is thus forced to rely on its [filesystem cache](https://docs.pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html#modular-caching-of-torchdynamo-torchinductor-and-triton) to restore compiled models. Depending on the model being compiled, this process takes from a few dozen seconds to a couple of minutes, which is way too much for practical GPU tasks in Spaces.\n\nThis is where **ahead-of-time (AoT) compilation** shines.\n\nWith AoT, we can export a compiled model once, save it, and later reload it instantly in any process, which is exactly what we need for ZeroGPU. This helps us reduce framework overhead and also eliminates cold-start timings typically incurred in just-in-time compilation.\n\nBut how can we do ahead-of-time compilation on ZeroGPU? Let’s dive in.\n\n## Ahead-of-time compilation on ZeroGPU\n\nLet's go back to our ZeroGPU base example and unpack what we need to enable AoT compilation. For the purpose of this demo, we will use the `black-forest-labs/FLUX.1-dev` model:\n\n```python\nimport gradio as gr\nimport spaces\nimport torch\nfrom diffusers import DiffusionPipeline\n\nMODEL_ID = 'black-forest-labs/FLUX.1-dev'\n\npipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)\npipe.to('cuda')\n\n@spaces.GPU\ndef generate(prompt):\n return pipe(prompt).images\n\ngr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\n> In the discussion below, we only compile the `transformer` component of `pipe` since, in these generative models, the transformer (or more generally, the denoiser) is the most computationally heavy component.\n\nCompiling a model ahead-of-time with PyTorch involves multiple steps:\n\n### 1. Getting example inputs\n\nRecall that we’re going to compile the model *ahead* of time. Therefore, we need to derive example inputs for the model. Note that these are the same kinds of inputs we expect to see during the actual runs. To capture those inputs, we will leverage the `spaces.aoti_capture` helper from the `spaces` package:\n\n```python\nwith spaces.aoti_capture(pipe.transformer) as call:\n pipe(\"arbitrary example prompt\")\n```\n\nWhen used as a context manager, `aoti_capture` intercepts the call to any callable (`pipe.transformer` in our case), prevents it from executing, captures the input arguments that would have been passed to it, and stores their values in `call.args` and `call.kwargs`.\n\n### 2. Exporting the model\n\nNow that we have example args and kwargs for our transformer component, we can export it to a PyTorch [`ExportedProgram`](https://docs.pytorch.org/docs/stable/export.html#torch.export.ExportedProgram) by using [`torch.export.export`](https://docs.pytorch.org/docs/stable/export.html#torch.export.export) utility:\n\n```python\nexported_transformer = torch.export.export(\n pipe.transformer,\n args=call.args,\n kwargs=call.kwargs,\n)\n```\n\nAn exported PyTorch program is a computation graph that represents the tensor computations along with the original model parameter values.\n\n### 3. Compiling the exported model\n\nOnce the model is exported, compiling it is pretty straightforward.\n\nA traditional AoT compilation in PyTorch often requires saving the model on disk so it can be later reloaded. In our case, we’ll leverage a helper function part of the `spaces` package: `spaces.aoti_compile`. It's a tiny wrapper around `torch._inductor.aot_compile` that manages saving and lazy-loading the model as needed. It's meant to be used like this:\n\n```python\ncompiled_transformer = spaces.aoti_compile(exported_transformer)\n```\n\nThis `compiled_transformer` is now an AoT-compiled binary ready to be used for inference.\n\n### 4. Using the compiled model in the pipeline\n\nNow we need to bind our compiled transformer to our original pipeline, i.e., the `pipeline`.\n\nA naive and almost working approach is to simply patch our pipeline like `pipe.transformer = compiled_transformer`. Unfortunately, this approach does not work because it deletes important attributes like `dtype`, `config`, etc. Only patching the `forward` method does not work well either because we are then keeping original model parameters in memory, often leading to OOM errors at runtime.\n\n`spaces` package provides a utility for this, too -- `spaces.aoti_apply`:\n\n```python\nspaces.aoti_apply(compiled_transformer, pipe.transformer)\n```\n\nEt voilà! It will take care of patching `pipe.transformer.forward` with our compiled model, as well as [cleaning old model parameters out of memory](https://pypi-browser.org/package/spaces/spaces-0.40.1-py3-none-any.whl/spaces/zero/torch/aoti.py#L87).\n\n### 5. Wrapping it all together\n\nTo perform the first three steps (intercepting input examples, exporting the model, and compiling it with PyTorch inductor), we need a real GPU. CUDA emulation that you get outside of `@spaces.GPU` function is not enough because compilation is truly hardware-dependent, for instance, relying on micro-benchmark runs to tune the generated code. This is why we need to wrap it all inside a `@spaces.GPU` function and then get our compiled model back to the root of our app. Starting from our original demo code, this gives:\n\n```diff\n import gradio as gr\n import spaces\n import torch\n from diffusers import DiffusionPipeline\n \n MODEL_ID = 'black-forest-labs/FLUX.1-dev'\n \n pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)\n pipe.to('cuda')\n \n+ @spaces.GPU(duration=1500) # maximum duration allowed during startup\n+ def compile_transformer():\n+ with spaces.aoti_capture(pipe.transformer) as call:\n+ pipe(\"arbitrary example prompt\")\n+ \n+ exported = torch.export.export(\n+ pipe.transformer,\n+ args=call.args,\n+ kwargs=call.kwargs,\n+ )\n+ return spaces.aoti_compile(exported)\n+ \n+ compiled_transformer = compile_transformer()\n+ spaces.aoti_apply(compiled_transformer, pipe.transformer)\n \n @spaces.GPU\n def generate(prompt):\n return pipe(prompt).images\n \n gr.Interface(generate, \"text\", \"gallery\").launch()\n```\n\nWith just a dozen lines of additional code, we’ve successfully made our demo quite faster (**1.7x** faster in the case of FLUX.1-dev).\n\nIf you want to learn more about AoT compilation, you can read PyTorch's [AOTInductor tutorial](https://docs.pytorch.org/tutorials/recipes/torch_export_aoti_python.html)\n\n## Gotchas\n\nNow that we have demonstrated the speedups one can realize under the constraints of operating with ZeroGPUs, we will discuss a few gotchas that came up while working with this setup.\n\n### Quantization\n\nAoT can be combined with quantization to deliver even greater speedups. For image and video generation, the FP8 post-training dynamic quantization schemes deliver good speed-quality trade-offs. However, FP8 requires a CUDA compute capability of at least 9.0 to work. Thankfully, for ZeroGPUs, since they’re based on H200s, we can already take advantage of the FP8 quantization schemes.\n\nTo enable FP8 quantization within our AoT compilation workflow, we can leverage the APIs provided by [`torchao`](https://github.com/pytorch/ao) like so:\n\n```diff\n+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig\n\n+ # Quantize the transformer just before the export step.\n+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())\n\nexported_transformer = torch.export.export(\n pipe.transformer,\n args=call.args,\n kwargs=call.kwargs,\n)\n```\n\n(You can find more details about TorchAO [here](https://docs.pytorch.org/ao/stable/index.html).)\n\nAnd we can then proceed with the rest of the steps as outlined above. Using quantization provides another **1.2x** of speedup.\n\n### Dynamic shapes\n\nImages and videos can come in different shapes and sizes. Hence, it’s important to also account for shape dynamism when performing AoT compilation. The primitives provided by `torch.export.export` make it easily configurable to provide what inputs should be treated accordingly for dynamic shapes, as shown below.\n\nFor the case of Flux.1-Dev transformer, changes in different image resolutions will affect two of its `forward` arguments:\n\n- `hidden_states`: The noisy input latents, which the transformer is supposed to denoise. It’s a 3D tensor, representing `batch_size, flattened_latent_dim, embed_dim`. When the batch size is fixed, it’s the `flattened_latent_dim` that will change for any changes made to image resolutions.\n- `img_ids`: A 2D array of encoded pixel coordinates having a shape of `height * width, 3`. In this case, we want to make `height * width` dynamic.\n\nWe start by defining a range in which we want to let the (latent) image resolutions vary. To derive these value ranges, we inspected the shapes of [`hidden_states`](https://github.com/huggingface/diffusers/blob/0ff1aa910cf3d87193af79ec1ae4487be542e872/src/diffusers/pipelines/flux/pipeline_flux.py#L920) in the pipeline with respect to varied image resolutions. The exact values are model-dependent and require manual inspection and some intuition. For Flux.1-Dev, we ended up with:\n\n```python\ntransformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)\n```\n\nWe then define a map of argument names and which dimensions in their input values we expect to be dynamic:\n\n```python\ntransformer_dynamic_shapes = {\n \"hidden_states\": {1: transformer_hidden_dim}, \n \"img_ids\": {0: transformer_hidden_dim},\n}\n```\n\nThen we need to make our dynamic shapes object replicate the structure of our example inputs. The inputs that do not need dynamic shapes must be set to `None`. This can be done very easily with PyTorch [tree\\_map](https://github.com/pytorch/pytorch/blob/2f0de0ff9361ca4f2b1e6f9edbc600b5fb6abcd6/torch/utils/_pytree.py#L1341-L1373) utility:\n\n```python\nfrom torch.utils._pytree import tree_map\n\ndynamic_shapes = tree_map(lambda v: None, call.kwargs)\ndynamic_shapes |= transformer_dynamic_shapes\n```\n\nNow, when performing the export step, we simply supply `transformer_dynamic_shapes` to `torch.export.export`:\n\n```python\nexported_transformer = torch.export.export(\n pipe.transformer,\n args=call.args,\n kwargs=call.kwargs,\n dynamic_shapes=dynamic_shapes,\n)\n```\n\n> Check out [this Space](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic) that shows how to use both quantization and dynamic shapes during the export step.\n\n### Multi-compile / shared weights\n\nDynamic shapes is sometimes not enough when dynamism is too important.\n\nThis is, for instance, the case with the Wan family of video generation models if you want your compiled model to generate different resolutions. One thing can be done in this case: compile one model per resolution while keeping the model parameters shared and dispatching the right one at runtime\n\nHere is a minimal example of this approach: [zerogpu-aoti-multi.py](https://gist.github.com/cbensimon/8dc0ffcd7ee024d91333f6df01907916). You can also see a fully working implementation of this paradigm in the [Wan 2.2 Space](https://huggingface.co/spaces/zerogpu-aoti/wan2-2-fp8da-aoti-faster/blob/main/optimization.py).\n\n### FlashAttention-3\n\nSince the ZeroGPU hardware and CUDA drivers are perfectly compatible with Flash-Attention 3 (FA3), we can use it in our ZeroGPU Spaces to speed things up even further. FA3 works with ahead-of-time compilation. So, this is ideal for our case.\n\nCompiling and building FA3 from source can take several minutes, and this process is hardware-dependent. As users, we wouldn’t want to lose precious ZeroGPU compute hours. This is where Hugging Face [`kernels` library](https://github.com/huggingface/kernels) comes to the rescue. It provides access to pre-built kernels that are compatible for a given hardware. For example, when we try to run:\n\n```python\nfrom kernels import get_kernel\n\nvllm_flash_attn3 = get_kernel(\"kernels-community/vllm-flash-attn3\")\n```\n\nIt tries to load a kernel from the [`kernels-community/vllm-flash-attn3`](https://huggingface.co/kernels-community/vllm-flash-attn3) repository, which is compatible with the current setup. Otherwise, it will error out due to incompatibility issues. Luckily for us, this works seamlessly on the ZeroGPU Spaces. This means we can leverage the power of FA3 on ZeroGPU, using the `kernels` library.\n\nHere is a [fully working example of an FA3 attention processor](https://gist.github.com/sayakpaul/ff715f979793d4d44beb68e5e08ee067#file-fa3_qwen-py) for the Qwen-Image model.\n\n### Regional compilation\n\nSo far, we have been compiling the full model. Depending on the model, full model compilation can lead to significantly long cold start times. Long cold start times make the development experience unpleasant.\n\nWe can also choose to compile *regions* within a model, significantly reducing the cold start times, while retaining almost all the benefits of full model compilation. Regional compilation becomes promising when a model has repeated blocks of computation. A standard language model, for example, has a number of identically structured Transformer blocks.\n\nIn our example, we can compile the repeated blocks of the Flux transformer ahead of time, and propagate the compiled graph to the remaining repeated blocks. The [Flux Transformer](https://github.com/huggingface/diffusers/blob/c2e5ece08bf22d249c62e964f91bc326cf9e3759/src/diffusers/models/transformers/transformer_flux.py) has two kinds of repeated blocks: `FluxTransformerBlock` and `FluxSingleTransformerBlock`.\n\nYou can check out [this Space](https://huggingface.co/spaces/cbensimon/FLUX.1-dev-fa3-aoti/tree/main) for a complete example.\n\n> 💡 For Flux.1-Dev, switching to regional compilation reduces the compilation time from *6 minutes* to just *30 seconds* while delivering identical speedups.\n\n### Use a compiled graph from the Hub\n\nOnce a model (or even a model block) is compiled ahead of time, we can serialize the compiled graph module as an artifact and reuse later. In the context of a ZeroGPU-powered demo on Spaces, this will significantly cut down the demo startup time by skipping the compilation time.\n\nTo keep the storage light, we can just save the compiled model graph without including any model parameters inside the artifact.\n\nCheck out [this collection](https://huggingface.co/collections/zerogpu-aoti/using-compiled-graph-from-the-hub-68c2afcc03de7609f9f91e35) that shows a full workflow of obtaining compiled model graph, pushing it to the Hub, and then using it to build a demo.\n\n## AoT compiled ZeroGPU Spaces demos\n\n### Speedup comparison\n\n- [FLUX.1-dev without AoTI](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-dev-base)\n- [FLUX.1-dev with AoTI and FA3](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-dev-fa3-aoti) (**1.75x** speedup)\n\n### Featured AoTI Spaces\n\n- [FLUX.1 Kontext](https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev)\n- [QwenImage Edit](https://huggingface.co/spaces/multimodalart/Qwen-Image-Edit-Fast)\n- [Wan 2.2](https://huggingface.co/spaces/zerogpu-aoti/wan2-2-fp8da-aoti-faster)\n\n### Regional compilation\n\n- [Regional compilation recipe](https://docs.pytorch.org/tutorials/recipes/regional_compilation.html)\n- [Regional compilation with AOT](https://docs.pytorch.org/tutorials/recipes/regional_aot.html)\n- [Native integration in Diffusers](https://huggingface.co/docs/diffusers/main/en/optimization/fp16)\n- [More performance numbers](https://pytorch.org/blog/torch-compile-and-diffusers-a-hands-on-guide-to-peak-performance/)\n\n## Conclusion\n\nZeroGPU within Hugging Face Spaces is a powerful feature that enables AI builders by providing access to powerful compute. In this post, we showed how users can benefit from PyTorch’s ahead-of-time compilation techniques to speed up their applications that leverage ZeroGPU.\n\nWe demonstrate speedups with Flux.1-Dev, but these techniques are not limited to just this model. Therefore, we encourage you to give these techniques a try and provide us with feedback in this [community discussion](https://huggingface.co/spaces/zerogpu-aoti/README/discussions/1).\n\n## Resources\n\n- Visit our [ZeroGPU-AOTI org on the Hub](https://huggingface.co/zerogpu-aoti) to refer to a collection of demos that leverage the techniques discussed in this post.\n- Browse `spaces.aoti_*` APIs [source code](https://pypi-browser.org/package/spaces/spaces-0.40.1-py3-none-any.whl/spaces/zero/torch/aoti.py) to learn more about the interface\n- Check out [Kernels Community org on the hub](https://huggingface.co/kernels-community)\n- Learn more about regional compilation from [here](https://huggingface.co/blog/pytorch.org/tutorials/recipes/regional_compilation.html)\n- Upgrade to [Pro](https://huggingface.co/pro) on Hugging Face to create your own ZeroGPU Spaces (and get 25 minutes of H200 usage every day)\n\n*Acknowledgements: Thanks to ChunTe Lee for creating an awesome thumbnail for this post. Thanks to Pedro and Vaibhav for providing feedback on the post. Thanks to Angela Yi from the PyTorch team for helping us with AOT guidance.*\n", "dateCreated": "2026-10-04T12:50:30.804479567+00:00", "date_created": "2025-09-02T00:00:00+00:00", "date_modified": "2025-09-02T00:00:00+00:00", "date_published": "2025-09-02T00:00:00+00:00", "description": "ZeroGPU lets anyone spin up powerful 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…", "external_url": "https://huggingface.co/blog/zerogpu-aoti", "git": { "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e", "path": "items/huggingface-co/2025/09/2025-09-02-make-your-zerogpu-spaces-go-brrr-with-ahead-of-time.md", "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/huggingface-co/2025/09/2025-09-02-make-your-zerogpu-spaces-go-brrr-with-ahead-of-time.md" }, "id": "urn:aggr:item:6eb2f07ddd7dce38a05ef35800ea1666c0fe0047", "isBasedOn": { "@id": "https://huggingface.co/blog/zerogpu-aoti" }, "mainEntity": { "@id": "https://huggingface.co/blog/zerogpu-aoti", "@type": "CreativeWork", "archivedAt": { "@id": "https://aggr.aymericbeaumet.com/items/huggingface-co/2025-09-02-make-your-zerogpu-spaces-go-brrr-with-ahead-of-time/#webpage" }, "dateModified": "2025-09-02T00:00:00+00:00", "datePublished": "2025-09-02T00:00:00+00:00", "headline": "Make your ZeroGPU Spaces go brrr with ahead-of-time compilation", "timeRequired": "PT13M", "url": "https://huggingface.co/blog/zerogpu-aoti", "wordCount": 2746 }, "name": "Archived snapshot: Make your ZeroGPU Spaces go brrr with ahead-of-time compilation", "preview": null, "reading_minutes": 13, "replicated_at": null, "source": "huggingface-co", "source_name": "Hugging Face - Blog", "summary": null, "tags": [], "temporalCoverage": "2026-10-04T12:50:30.804479567+00:00", "title": "Make your ZeroGPU Spaces go brrr with ahead-of-time compilation", "url": "https://aggr.aymericbeaumet.com/items/huggingface-co/2025-09-02-make-your-zerogpu-spaces-go-brrr-with-ahead-of-time/", "word_count": 2746 }