@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f16bfbe5bc
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 11:30 UTC

[Documentation][torch.compile] Add documentation for torch.compile + multimodal encoders (#31627)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com>

Lucas Kabela committed Jan 8, 2026 at 11:33 UTC f16bfbe5bc8f9045e7bb38d72ab35764bab22e44
1 file changed +111
docs/design/torch_compile_multimodal.md new
+111
@@ -0,0 +1,111 @@
1 +# torch.compile with Multimodal Encoders
2 +
3 +`torch.compile` can now be applied to multimodal encoders and miscellaneous nn modules in vLLM, including vision-language models like LLaMA 4, Qwen-VL,
4 +and similar encoder-based architectures.
5 +
6 +This document covers the basics of how the `torch.compile` integration works for multimodal encoders in vLLM, as well as how to apply the decorator
7 +to new models to improve performance.
8 +
9 +!!! note
10 + For general information about `torch.compile` integration in vLLM, see the [torch.compile design document](./torch_compile.md).
11 +
12 +## Overview
13 +
14 +We have recently enabled the `@supports_torch_compile` decorator to work for multiple nn module components within a model type; this enables
15 +turning compile on for multimodal encoders, bringing performance improvements to additional components of the stack.
16 +
17 +When applied to the vision block of [`Qwen2_5_vl`](https://github.com/vllm-project/vllm/pull/23207) we observe ~4.5% e2e perf improvements with
18 +some increase in compilation time
19 +
20 +This feature is off by default, but can be enabled by setting `compile_mm_encoder: true` in the compilation config when models have the
21 +`@supports_torch_compile` decorator.
22 +
23 +## How Compilation Works for Multimodal Components
24 +
25 +### APIs for Enablement
26 +
27 +To compile a multimodal component such as an encoder, we follow the same mechanism as the LLM text backbone, with a few additional scaffoldings:
28 +
29 +1. The `@supports_torch_compile` decorator should include `enable_if=should_torch_compile_mm_vit`. This will gate the compilation behind our
30 +`compile_mm_encoder` configuration
31 +
32 +2. `with set_model_tag("<component_name>", is_encoder=True)` context manager should be used around the nn.Module's instantiation. Since torch.compile
33 +relies on caching artifacts to reduce start time, we must properly propagate the `<component_name>` information to the cache in order to avoid collisions
34 +with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
35 +components (see Compile Range Integration).
36 +
37 +3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
38 +for torch.compile integration.
39 +
40 +### CompilationConfig
41 +
42 +With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend
43 +this for more configuration in the future.
44 +
45 +## Applying torch.compile to a New Multimodal Model/Component
46 +
47 +To apply `supports_torch_compile` to a new general nn.Module, we advise following the same steps in [`debug_vllm_compile`](./debug_vllm_compile.md); this includes:
48 +
49 +1. Applying `supports_torch_compile` on initially small modules (such as basic MLP layers), then raising to more general modules until one reaches a good performance
50 +tradeoff
51 +
52 +2. Leveraging [`tlparse`](https://github.com/meta-pytorch/tlparse) to identify and eliminate the source of recompiles and graph breaks
53 +
54 +3. Using `dynamic_arg_dims` and proper `dynamic_shapes_config` to handle dynamism.
55 +
56 +### Common pitfalls
57 +
58 +## VllmBackend Feature Support
59 +
60 +### Compile ranges
61 +
62 +The torch.compile integration will try to rely on max_batch_size to infer compilation ranges for dynamic shapes; however, for modules used in the encoder, this
63 +shape can be difficult to infer due to the unspecified range of shapes the encoder may see as input. Therefore, we rely on `is_encoder=True` in the `set_model_tag`
64 +to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).
65 +
66 +!!! note
67 + We may seek to tighten this range for better performance in the future
68 +
69 +### Cudagraphs
70 +
71 +We have not yet explored compilation for multimodal encoders with CUDAGraph integration; behavior is currently unspecified.
72 +
73 +## Troubleshooting
74 +
75 +### Graph Breaks in Vision Encoders
76 +
77 +Some vision encoder operations may cause graph breaks. To identify them:
78 +
79 +```bash
80 +TORCH_LOGS="+dynamo" vllm serve <MODEL>
81 +```
82 +
83 +Common causes of graph breaks in multimodal models:
84 +
85 +- **Dynamic image sizes**: Use `dynamic_shapes_config` to handle variable resolutions
86 +- **Untraceable operations**: Some operations (such as to_list) may not be supported by Dynamo
87 +- **Conditional processing**: Data-dependent branching based on image properties
88 +
89 +### Compilation Errors
90 +
91 +If compilation fails for a multimodal model:
92 +
93 +1. **Disable and test**: First verify the model works without compilation:
94 + ```bash
95 + VLLM_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'
96 + ```
97 +
98 +2. **Check logs**: Enable debug logging to see compilation details:
99 + ```bash
100 + VLLM_LOGGING_LEVEL=DEBUG vllm serve <model> --compilation-config='{"compile_mm_encoder":"true"}'
101 + ```
102 +
103 +3. **Report issues**: If you find a bug, [open an issue on GitHub](https://github.com/vllm-project/vllm/issues/new/choose)
104 +
105 +## See Also
106 +
107 +- [torch.compile Integration](./torch_compile.md) - Core design document
108 +- [Debugging torch.compile](./debug_vllm_compile.md) - Detailed debugging guide
109 +- [Multimodal Inputs](../features/multimodal_inputs.md) - How to pass multimodal data
110 +- [Disaggregated Encoder](../features/disagg_encoder.md) - Scaling vision encoders
111 +- [Supported Multimodal Models](../models/supported_models.md#list-of-multimodal-language-models) - Model compatibility