Model compilation transforms groups of operations in a model’s computation graph into optimized kernels, reducing Python overhead during repeated inference calls. With PyTorch, apply it as `model_compiled = torch.compile(model)`, then use the compiled model in the existing generation function. section 2.9
Use compilation as a measured inference optimization:
- Compile the loaded model, optionally using `mode="max-autotune"`. section 2.9
- Run a warm-up generation before timing, because the first execution includes compilation and optimization overhead. section 2.9
- Repeat generation several times and compare tokens per second with the uncompiled implementation. section 2.9
- Combine compilation with KV caching when generation repeatedly evaluates the same model code; the measured improvement is substantially larger. section 2.9 section 2.10
Measured generation speed
On a Mac Mini M4 CPU, regular generation increased from about 5 to 6 tokens per second after compilation. With a KV cache, compilation increased throughput from 29 to 68 tokens per second. On an NVIDIA H100, regular generation rose from 51 to 164 tokens per second, while GPU memory increased from 1.55 GB to 1.81 GB. section 2.9
Figure 2.17 An overview of the four key stages in developing a reasoning model in this book. This section builds on pretrained LLM and the basic text generation function we coded earlier, including KV caching, and adds model compilation to speed up the execution speed even further.
Use the workflow overview to place model compilation after pretrained-model loading, basic generation, and KV caching. section 2.9
