$
~/blog
~/blog/where-torch.compile-meets-vllm

Where torch.compile meets vLLM

After looking at vLLM’s scheduler and KV cache, I wanted to follow what happens inside model execution. Where does torch.compile enter the picture, and how much work does vLLM add around it?

This talk follows that integration one layer at a time. The original slides are below, with short explanations and full-size images available on click. Some class names and defaults change between releases, so the diagrams are most useful as a map of the design.

1. A closer look at vLLM compilation

Slide 1: A closer look at vLLM compilation

vLLM uses PyTorch’s compiler machinery, but it also manages graph boundaries, cache keys, supported shapes, and CUDA graph execution around it. The interesting part is how those pieces fit a serving workload. A model that runs repeatedly with changing batches needs more than a one-time call to a compiler.

2. The joke behind the title

Slide 2: The joke behind the title

The meme captures the starting observation: follow vLLM’s compilation path far enough and you’ll meet torch.compile and its internals. There is still substantial vLLM-specific machinery along the way. That machinery adapts graph capture and execution to the constraints of an inference engine.

3. What is vLLM?

Slide 3: What is vLLM?

vLLM is a library and serving system for LLM inference. It coordinates incoming requests, schedules tokens, manages KV state, and executes models on supported devices. Compilation lives inside that larger system, so understanding its location helps keep model optimization separate from request orchestration.

4. Compilation is one of several optimizations

Slide 4: Compilation is one of several optimizations

Paged KV storage, efficient kernels, chunked prefill, and speculative decoding all address different costs. Compiler-generated work fits alongside them. If a workload is limited by KV capacity or communication, making one pointwise operation faster may have little effect on the whole service. The surrounding system still determines how much useful work reaches the GPU.

5. Start with the architecture

Slide 5: Start with the architecture

Before following the compiler, locate the engine, scheduler, executor, worker, and model runner. The compiler works much closer to the model’s tensor computation than to the API server. Keeping that boundary in mind prevents a common confusion: compiling a model doesn’t mean compiling the entire request-handling loop.

6. Requests flow into model execution

Slide 6: Requests flow into model execution

Input processing feeds the engine core. Scheduling and KV allocation determine the work that the executor sends to workers, and results return through output processing. The diagram shows where these responsibilities connect. Compiler optimizations help execute the chosen model work; they don’t replace the scheduler’s decisions about which requests should run.

7. From the engine to the worker

Slide 7: From the engine to the worker

The call graph breaks the system into initialization and execution responsibilities. Workers prepare devices and model resources, while the engine core owns the repeated scheduling cycle. The model runner prepares inputs and attention metadata before invoking the model. That is the point where the serving system approaches the compiled computation.

8. What torch.compile contributes

Slide 8: What torch.compile contributes

torch.compile can capture supported regions of PyTorch code and hand them to an optimizing backend. Dynamo performs capture, while a backend such as Inductor produces executable work. vLLM uses these capabilities as building blocks, with extra control over the shapes, transformations, and runtime paths relevant to serving.

9. The compiler pipeline refresher

Slide 9: The compiler pipeline refresher

The diagram follows Dynamo, AOTAutograd, and Inductor. For inference, AOTAutograd’s machinery still helps prepare and transform the graph even though a backward graph isn’t needed. Inductor then lowers supported operations and generates code or external calls. This is the compiler path that vLLM customizes around model execution.

10. The exact boundary in the serving stack

Slide 10: The exact boundary in the serving stack

Follow the stack down from the request to GPUModelRunner and the model’s forward call. Tokenization, queuing, and much of input preparation happen outside the compiled region. The repeated tensor computation is where graph compilation and CUDA graph replay can reduce overhead. This boundary is one of the main design decisions in the integration.

11. The layers of the compilation pipeline

Slide 11: The layers of the compilation pipeline

The wrapper prepares a model for capture, VllmBackend manages the graph, piecewise machinery handles supported regions and shapes, and CUDA graph wrappers manage capture and replay. These are separate jobs. In particular, compiling a region and recording its device execution in a CUDA graph are different transformations with different reuse conditions.

12. The decorator and wrapper

Slide 12: The decorator and wrapper

support_torch_compile connects eligible model classes to vLLM’s compilation setup, including dynamic-dimension information and a lazy first-call path. Later calls can use prepared execution paths. The slide’s guard-dropping shorthand relies on assumptions enforced by the serving system; it isn’t a general suggestion to remove correctness checks from arbitrary PyTorch programs.

13. What the vLLM backend does

Slide 13: What the vLLM backend does

The backend receives an FX graph, manages compilation caching, and can split around configured operations before compiling supported pieces. Attention is an important example in the piecewise design shown here. The selected attention backend still executes that work. The VllmBackend source also shows why full-graph and partitioned configurations shouldn’t be collapsed into one universal path.

14. Dispatching by shape

Slide 14: Dispatching by shape

One compiled graph can cover a range of shapes, while selected sizes can have more specialized implementations. At runtime, the piecewise backend chooses the appropriate callable. The diagram labels the dimension as batch size, but token count is also central in serving, especially when prefill contributes several tokens per request. The concrete shape contract matters more than the shorthand label.

15. The compiler interface

Slide 15: The compiler interface

The compiler interface separates vLLM’s orchestration from backend-specific compile and load behavior. The slide includes multiple Inductor adapters and an eager path for comparison or debugging. This means the integration isn’t tied to one identical backend path in every configuration. The checked implementation is in compiler_interface.py.

16. Passes that understand the workload

Slide 16: Passes that understand the workload

vLLM can apply transformations around common inference patterns, such as normalization followed by quantization. Combining compatible work can reduce intermediate memory traffic and launch overhead. The exact passes and their order depend on configuration and implementation version. A fusion still needs to preserve the operation’s numerical and mutation behavior; matching a familiar name is not enough.

17. CUDA graph capture and replay

Slide 17: CUDA graph capture and replay

After code has been compiled or otherwise prepared, a CUDA graph can record a compatible execution sequence. Replaying it reduces repeated CPU launch work. Input addresses, shapes, and other capture assumptions still matter, and replay is not literally zero overhead. vLLM’s compilation debugging guide usefully separates disabling compilation from disabling CUDA graphs.

18. References and source navigation

Slide 18: References and source navigation

The vLLM compilation source is the best companion to these diagrams. The original slide also links to the Anatomy of vLLM article and my earlier architecture talk. For the foundations behind the KV layout, see the PagedAttention paper.

19. The next question: what should the graph represent?

Slide 19: The next question: what should the graph represent?

Once the integration makes sense, another problem appears: the same high-level operation can show up as different graph patterns for different kernels. That makes compiler passes harder to maintain. My vLLM IR talk picks up that question and looks at separating an operation’s meaning from its implementation.

cd ..
$ grep -r