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 6, 15:45 UTC
[torch.compile][BE] Modify cudagraph callable to check for is_forward_context_set (#36288)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Mar 16, 2026 at 12:42 UTC
714c6e0eab76a4fb1394089d848ecfe46408b9c9
4 files changed
+27
-29
docs/design/torch_compile_multimodal.md
-3
@@ -34,9 +34,6 @@ relies on caching artifacts to reduce start time, we must properly propagate the
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
-
37
### CompilationConfig
38
39
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
vllm/compilation/cuda_graph.py
+11
-1
@@ -16,7 +16,11 @@ from vllm.compilation.counter import compilation_counter
16
from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
17
from vllm.config import CUDAGraphMode, VllmConfig
18
from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
19
-from vllm.forward_context import BatchDescriptor, get_forward_context
19
+from vllm.forward_context import (
20
+ BatchDescriptor,
21
+ get_forward_context,
22
+ is_forward_context_available,
23
+)
24
from vllm.logger import init_logger
25
from vllm.model_executor.offloader.base import get_offloader
26
from vllm.platforms import current_platform
@@ -224,6 +228,12 @@ class CUDAGraphWrapper:
228
self.concrete_cudagraph_entries.clear()
229
230
def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
231
+ if not is_forward_context_available():
232
+ # No forward context means we are outside the normal
233
+ # inference path (e.g. a vision encoder forward pass).
234
+ # Just run the underlying function without cudagraphs.
235
+ return self.runnable(*args, **kwargs)
236
+
237
forward_context = get_forward_context()
238
batch_descriptor = forward_context.batch_descriptor
239
cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
vllm/model_executor/models/mllama4.py
+1
-5
@@ -38,7 +38,6 @@ from vllm.compilation.decorators import (
38
from vllm.config import VllmConfig, set_current_vllm_config
39
from vllm.config.multimodal import BaseDummyOptions
40
from vllm.distributed import get_tensor_model_parallel_world_size
41
-from vllm.forward_context import set_forward_context
41
from vllm.model_executor.layers.attention import MMEncoderAttention
42
from vllm.model_executor.layers.fused_moe import FusedMoE
43
from vllm.model_executor.layers.linear import (
@@ -872,10 +871,7 @@ class Llama4ForConditionalGeneration(
871
if image_input is None:
872
return []
873
875
- with (
876
- set_forward_context(None, self.vllm_config),
877
- ):
878
- return self._process_image_input(image_input)
874
+ return self._process_image_input(image_input)
875
876
def forward(
877
self,
vllm/model_executor/models/qwen2_5_vl.py
+15
-20
@@ -49,7 +49,6 @@ from vllm.compilation.decorators import (
49
from vllm.config import VllmConfig
50
from vllm.distributed import parallel_state
51
from vllm.distributed import utils as dist_utils
52
-from vllm.forward_context import set_forward_context
52
from vllm.logger import init_logger
53
from vllm.model_executor.layers.activation import get_act_and_mul_fn
54
from vllm.model_executor.layers.attention import MMEncoderAttention
@@ -1207,13 +1206,12 @@ class Qwen2_5_VLForConditionalGeneration(
1206
image_embeds = image_input["image_embeds"].type(self.visual.dtype)
1207
else:
1208
pixel_values = image_input["pixel_values"]
1210
- with set_forward_context(None, self.vllm_config):
1211
- if self.use_data_parallel:
1212
- return run_dp_sharded_mrope_vision_model(
1213
- self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
1214
- )
1215
- else:
1216
- image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
1209
+ if self.use_data_parallel:
1210
+ return run_dp_sharded_mrope_vision_model(
1211
+ self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
1212
+ )
1213
+ else:
1214
+ image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
1215
1216
# Split concatenated embeddings for each image item.
1217
merge_size = self.visual.spatial_merge_size
@@ -1262,18 +1260,15 @@ class Qwen2_5_VLForConditionalGeneration(
1260
video_embeds = video_input["video_embeds"].type(self.visual.dtype)
1261
else:
1262
pixel_values_videos = video_input["pixel_values_videos"]
1265
- with set_forward_context(None, self.vllm_config):
1266
- if self.use_data_parallel:
1267
- return run_dp_sharded_mrope_vision_model(
1268
- self.visual,
1269
- pixel_values_videos,
1270
- grid_thw_list,
1271
- rope_type="rope_3d",
1272
- )
1273
- else:
1274
- video_embeds = self.visual(
1275
- pixel_values_videos, grid_thw=grid_thw_list
1276
- )
1263
+ if self.use_data_parallel:
1264
+ return run_dp_sharded_mrope_vision_model(
1265
+ self.visual,
1266
+ pixel_values_videos,
1267
+ grid_thw_list,
1268
+ rope_type="rope_3d",
1269
+ )
1270
+ else:
1271
+ video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list)
1272
1273
# Split concatenated embeddings for each video item.
1274
merge_size = self.visual.spatial_merge_size