[Platform] Make forward context manager pluggable for other device

vllm/model_executor/models/qwen2_5_vl.py

@@ -51,7 +51,6 @@
 from vllm.config import VllmConfig
 from vllm.distributed import parallel_state
 from vllm.distributed import utils as dist_utils
-from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_and_mul_fn
 from vllm.model_executor.layers.conv import Conv3dLayer
@@ -1247,6 +1246,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.language_model.make_empty_intermediate_tensors
         )
 
+        from vllm.platforms import current_platform
+
+        self.set_forward_context = current_platform.get_forward_context_manager()
+
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         self.language_model.model.aux_hidden_state_layers = layers
 
@@ -1316,7 +1319,7 @@ def _process_image_input(
             image_embeds = image_input["image_embeds"].type(self.visual.dtype)
         else:
             pixel_values = image_input["pixel_values"]
-            with set_forward_context(None, self.vllm_config):
+            with self.set_forward_context(None, self.vllm_config):
                 if self.use_data_parallel:
                     return run_dp_sharded_mrope_vision_model(
                         self.visual, pixel_values, grid_thw_list, rope_type="rope_3d"
@@ -1371,7 +1374,7 @@ def _process_video_input(
             video_embeds = video_input["video_embeds"].type(self.visual.dtype)
         else:
             pixel_values_videos = video_input["pixel_values_videos"]
-            with set_forward_context(None, self.vllm_config):
+            with self.set_forward_context(None, self.vllm_config):
                 if self.use_data_parallel:
                     return run_dp_sharded_mrope_vision_model(
                         self.visual,

vllm/platforms/interface.py

@@ -653,6 +653,15 @@ def check_max_model_len(cls, max_model_len: int) -> int:
         """
         return max_model_len
 
+    @classmethod
+    def get_forward_context_manager(cls):
+        """
+        Returns forward context manager for the current platform.
+        """
+        from vllm.forward_context import set_forward_context
+
+        return set_forward_context
+
 
 class UnspecifiedPlatform(Platform):
     _enum = PlatformEnum.UNSPECIFIED