vit and llm inference at a single process

shihaobai · shihaobai · commit d498aafb4844 · 2025-07-01T19:21:13.000+08:00
diff --git a/lightllm/common/image_cache_manager.py b/lightllm/common/image_cache_manager.py
@@ -0,0 +1,74 @@
+from collections import OrderedDict
+from lightllm.utils.dist_utils import get_current_device_id
+
+
+class ImageCacheManager:
+    def __init__(self):
+        """
+        Initialize the image cache manager with a simple GPU cache and an LRU CPU cache.
+        """
+        self._gpu_cache = dict()
+        self._cpu_cache = OrderedDict()
+
+    def set_max_size(self, max_size: int):
+        """
+        Set the maximum number of items to keep in the CPU cache.
+        :param max_size: Maximum number of items to keep in the CPU cache.
+        """
+        if max_size <= 0:
+            raise ValueError("max_size must be greater than 0")
+        self._max_size = max_size
+
+    def set_embed(self, uuid, embed):
+        """
+        Store the embedding for the given uuid in the GPU cache.
+        :param uuid: Unique identifier for the image
+        :param embed: Embedding vector for the image (on GPU)
+        """
+        self._gpu_cache[uuid] = embed
+
+    def get_embed(self, uuid):
+        """
+        Retrieve the embedding for the given uuid. Prefer GPU cache,
+        otherwise return CPU cache and move to GPU (simulate .cuda()).
+        :param uuid: Unique identifier for the image
+        :return: Embedding vector (on GPU if possible, else move from CPU to GPU)
+        """
+        if uuid in self._gpu_cache:
+            return self._gpu_cache[uuid]
+        elif uuid in self._cpu_cache:
+            self._cpu_cache.move_to_end(uuid)
+            embed = self._cpu_cache[uuid].cuda(get_current_device_id())
+            return embed
+        return None
+
+    def query_embed(self, uuid):
+        """
+        Query if the embedding for the given uuid is in the cache.
+        :param uuid: Unique identifier for the image
+        :return: True if the embedding is in the cache, False otherwise
+        """
+        return uuid in self._gpu_cache or uuid in self._cpu_cache
+
+    def filter(self, uuid_list):
+        """
+        Given a list of uuids, move their embeddings from GPU cache to CPU cache if present,
+        and return a dict of those found in the cache and their embeddings (on CPU).
+        :param uuid_list: List of uuids
+        """
+        for uuid in uuid_list:
+            if uuid in self._gpu_cache:
+                embed_cpu = self._gpu_cache[uuid].cpu(non_blocking=True)
+                # Move to CPU cache and remove from GPU cache
+                self._gpu_cache.pop(uuid)
+                if uuid in self._cpu_cache:
+                    self._cpu_cache.move_to_end(uuid)
+                self._cpu_cache[uuid] = embed_cpu
+                if len(self._cpu_cache) > self._max_size:
+                    self._cpu_cache.popitem(last=False)
+            elif uuid in self._cpu_cache:
+                self._cpu_cache.move_to_end(uuid)
+        return
+
+
+image_cache_manager = ImageCacheManager()
diff --git a/lightllm/models/internvl/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/internvl/layer_weights/pre_and_post_layer_weight.py
@@ -3,6 +3,9 @@
 from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
 
 from lightllm.models.internlm2.layer_weights.pre_and_post_layer_weight import Internlm2PreAndPostLayerWeight
+from lightllm.models.vit.model import VisionTransformer
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.common.image_cache_manager import image_cache_manager
 
 
 # add key: language_model.xxx -> xxx
@@ -15,9 +18,45 @@ def rename_weight_keys(weights):
             weights[k[len(prefix) :]] = weights[k]
 
 
+class InternVLPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
+    def __init__(self, data_type, network_config, mode):
+        super().__init__(data_type, network_config, mode)
+        # if we don't assign an extra process for visual model, we need initialize the image cache manager here
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            kvargs = {
+                "weight_dir": get_env_start_args().model_dir,
+                "data_type": self.data_type_,
+                "quant_type": get_env_start_args().vit_quant_type,
+                "quant_cfg": get_env_start_args().vit_quant_cfg,
+                "max_batch_size": get_env_start_args().visual_infer_batch_size,
+            }
+            self.visual_model = VisionTransformer(
+                kvargs=kvargs,
+            )
+            image_cache_manager.set_max_size(get_env_start_args().cache_capacity * 2)
+        return
+
+    def load_hf_weights(self, weights):
+        rename_weight_keys(weights)
+        super().load_hf_weights(weights)
+
+
 class InternVLPhi3PreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
     def __init__(self, data_type, network_config, mode):
         super().__init__(data_type, network_config, mode)
+        # if we don't assign an extra process for visual model, we need initialize the image cache manager here
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            kvargs = {
+                "weight_dir": get_env_start_args().model_dir,
+                "data_type": self.data_type_,
+                "quant_type": get_env_start_args().vit_quant_type,
+                "quant_cfg": get_env_start_args().vit_quant_cfg,
+                "max_batch_size": get_env_start_args().visual_infer_batch_size,
+            }
+            self.visual_model = VisionTransformer(
+                kvargs=kvargs,
+            )
+            image_cache_manager.set_max_size(get_env_start_args().cache_capacity * 2)
         return
 
     def load_hf_weights(self, weights):
@@ -29,6 +68,19 @@ def load_hf_weights(self, weights):
 class InternVLInternlm2PreAndPostLayerWeight(Internlm2PreAndPostLayerWeight):
     def __init__(self, data_type, network_config, mode):
         super().__init__(data_type, network_config, mode)
+        # if we don't assign an extra process for visual model, we need initialize the image cache manager here
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            kvargs = {
+                "weight_dir": get_env_start_args().model_dir,
+                "data_type": self.data_type_,
+                "quant_type": get_env_start_args().vit_quant_type,
+                "quant_cfg": get_env_start_args().vit_quant_cfg,
+                "max_batch_size": get_env_start_args().visual_infer_batch_size,
+            }
+            self.visual_model = VisionTransformer(
+                kvargs=kvargs,
+            )
+            image_cache_manager.set_max_size(get_env_start_args().cache_capacity * 2)
         return
 
     def load_hf_weights(self, weights):
@@ -40,6 +92,19 @@ def load_hf_weights(self, weights):
 class InternVLLlamaPreAndPostLayerWeight(LlamaPreAndPostLayerWeight):
     def __init__(self, data_type, network_config, mode):
         super().__init__(data_type, network_config, mode)
+        # if we don't assign an extra process for visual model, we need initialize the image cache manager here
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            kvargs = {
+                "weight_dir": get_env_start_args().model_dir,
+                "data_type": self.data_type_,
+                "quant_type": get_env_start_args().vit_quant_type,
+                "quant_cfg": get_env_start_args().vit_quant_cfg,
+                "max_batch_size": get_env_start_args().visual_infer_batch_size,
+            }
+            self.visual_model = VisionTransformer(
+                kvargs=kvargs,
+            )
+            image_cache_manager.set_max_size(get_env_start_args().cache_capacity * 2)
         return
 
     def load_hf_weights(self, weights):
diff --git a/lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py b/lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py
@@ -6,7 +6,9 @@
 
 from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer
 from lightllm.utils.infer_utils import mark_cost_time
+from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.server.embed_cache.utils import bytes2tensor, read_shm, get_shm_name_embed
+from lightllm.common.image_cache_manager import image_cache_manager
 from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb
 from lightllm.distributed.communication_op import all_reduce
 
@@ -29,8 +31,22 @@
 class LlamaMultimodalPreLayerInfer(LlamaPreLayerInfer):
     def __init__(self, network_config, mode):
         super().__init__(network_config, mode)
+        self.disable_extra_process_for_multimodal = get_env_start_args().disable_extra_process_for_multimodal
         return
 
+    def _infer_image_embeds(self, infer_state, layer_weight):
+        if not self.disable_extra_process_for_multimodal:
+            return
+        infer_images = []
+        for _, p in enumerate(infer_state.multimodal_params):
+            for img in p["images"] + p["audios"]:
+                if not image_cache_manager.query_embed(img["uuid"]):
+                    infer_images.append(img)
+        if len(infer_images) > 0:
+            img_embeds, uuids, valid_ids = layer_weight.visual_model.encode(infer_images)
+            for uuid, valid_id in zip(uuids, valid_ids):
+                image_cache_manager.set_embed(uuid, img_embeds[valid_id[0] : valid_id[1]])
+
     def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
 
         img_weight = []
@@ -42,14 +58,20 @@ def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_wei
         device = layer_weight.wte_weight_.device
         dtype = layer_weight.wte_weight_.dtype
         hidden_size = layer_weight.wte_weight_.shape[1]
+        self._infer_image_embeds(infer_state, layer_weight)
         for batch_id, p in enumerate(infer_state.multimodal_params):
             for img in p["images"] + p["audios"]:
                 # skip the same image
                 if img["token_id"] in img_start_token_ids:
                     continue
                 # pull the img_embeds by uid from shm
-                data = read_shm(get_shm_name_embed(img["uuid"]))
-                img_weight.append(bytes2tensor(data).cuda().reshape(img["token_num"], -1))
+                if self.disable_extra_process_for_multimodal:
+                    img_embed = image_cache_manager.get_embed(img["uuid"])
+                    img_weight.append(img_embed.reshape(img["token_num"], -1))
+                    print(img_weight[-1].shape)
+                else:
+                    data = read_shm(get_shm_name_embed(img["uuid"]))
+                    img_weight.append(bytes2tensor(data).cuda().reshape(img["token_num"], -1))
                 img_start_token_ids.append(img["token_id"])
                 img_token_lens.append(img["token_num"])
                 img_start_locs.append(img_start_loc)
diff --git a/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/vit/layer_weights/pre_and_post_layer_weight.py
@@ -3,7 +3,12 @@
 import numpy as np
 import torch.nn.functional as F
 from lightllm.common.basemodel import PreAndPostLayerWeight
-from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.dist_utils import (
+    get_current_device_id,
+    get_global_rank,
+    get_global_world_size,
+)
+from lightllm.utils.envs_utils import get_env_start_args
 
 
 class ViTPreAndPostLayerWeight(PreAndPostLayerWeight):
@@ -13,6 +18,10 @@ def __init__(self, data_type, network_config, mode):
         self.image_size = self.network_config_["image_size"]
         self.patch_size = self.network_config_["patch_size"]
         self.llm_hidden_size = self.network_config_["llm_hidden_size"]
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            self.tp_world_size_ = get_global_world_size()
+            self.tp_rank_ = get_global_rank()
+
         return
 
     def _cuda(self, cpu_tensor):
diff --git a/lightllm/models/vit/layer_weights/transformer_layer_weight.py b/lightllm/models/vit/layer_weights/transformer_layer_weight.py
@@ -11,12 +11,20 @@
     MultiROWMMWeight,
     TpNormWeight,
 )
-from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.dist_utils import (
+    get_current_device_id,
+    get_global_rank,
+    get_global_world_size,
+)
+from lightllm.utils.envs_utils import get_env_start_args
 
 
 class ViTTransformerLayerWeight(TransformerLayerWeight):
     def __init__(self, layer_num, data_type, network_config, mode=[], quant_cfg=None):
         super().__init__(layer_num, data_type, network_config, mode, quant_cfg)
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            self.tp_world_size_ = get_global_world_size()
+            self.tp_rank_ = get_global_rank()
         return
 
     def _cuda(self, cpu_tensor):
diff --git a/lightllm/models/vit/model.py b/lightllm/models/vit/model.py
@@ -18,7 +18,8 @@
 from io import BytesIO
 from rpyc.utils.classic import obtain
 from lightllm.common.quantization import Quantcfg
-from lightllm.utils.dist_utils import get_dp_world_size
+from lightllm.utils.dist_utils import get_dp_world_size, get_global_world_size
+from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
 
 
@@ -37,7 +38,11 @@ class VisionTransformer:
     post_layer_infer_class = ViTPostLayerInfer
 
     def __init__(self, kvargs):
-        self.tp_world_size_ = get_dp_world_size()
+        if get_env_start_args().disable_extra_process_for_multimodal:
+            # if we don't assign an extra process for visual model, the visual model uses tensor parallel by default.
+            self.tp_world_size_ = get_global_world_size()
+        else:
+            self.tp_world_size_ = get_dp_world_size()
         self.weight_dir_ = kvargs["weight_dir"]
         self.load_way = kvargs.get("load_way", "HF")
         self.mode = [m.replace("int4weight", "w4a16").replace("int8weight", "w8a16") for m in kvargs.get("mode", [])]
@@ -150,6 +155,8 @@ def _init_infer_layer(self):
         return
 
     def _init_datatype(self):
+        if isinstance(self.data_type, torch.dtype):
+            return
         if self.data_type in ["fp16", "float16"]:
             self.data_type = torch.float16
         elif self.data_type in ["bf16", "bfloat16"]:
@@ -161,12 +168,14 @@ def _init_datatype(self):
 
     @torch.no_grad()
     def forward(self, pixel_values):
-        g_cache_manager.cache_env_in()
+        if not get_env_start_args().disable_extra_process_for_multimodal:
+            g_cache_manager.cache_env_in()
         input_embs = self.pre_infer.forward(pixel_values, self.pre_post_weight)
         for i in range(self.layers_num + self.select_layer + 1):
             input_embs = self.layers_infer[i].forward(input_embs, self.trans_layers_weight[i])
         input_embs = self.post_infer.forward(input_embs[:, 1:, :], self.pre_post_weight)
-        g_cache_manager.cache_env_out()
+        if not get_env_start_args().disable_extra_process_for_multimodal:
+            g_cache_manager.cache_env_out()
         return input_embs
 
     @torch.no_grad()
@@ -182,6 +191,12 @@ def encode(self, images: List[ImageItem]):
                 image_data = Image.open(BytesIO(image_data))
                 t = self.load_image_func(image_data, max_num=img.extra_params["image_patch_max_num"])
                 img_tensors.append(t)
+            elif isinstance(img, dict):
+                uuids.append(img["uuid"])
+                image_data = read_shm(get_shm_name_data(img["uuid"]))
+                image_data = Image.open(BytesIO(image_data))
+                t = self.load_image_func(image_data, max_num=img["extra_params"]["image_patch_max_num"])
+                img_tensors.append(t)
             else:
                 raise Exception("Unsupport input types: {} for {}".format(type(img), img))
 
diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py
@@ -233,6 +233,11 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--enable_multimodal", action="store_true", help="Whether or not to allow to load additional visual models."
     )
+    parser.add_argument(
+        "--disable_extra_process_for_multimodal",
+        action="store_true",
+        help="Whether or not to disable extra process for multimodal.",
+    )
     parser.add_argument(
         "--enable_multimodal_audio",
         action="store_true",
diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py
@@ -243,7 +243,7 @@ def normal_or_p_d_start(args):
             ],
             start_args=[(cache_port, args)],
         )
-        if args.enable_multimodal_audio:
+        if args.enable_multimodal_audio and not args.disable_extra_process_for_multimodal:
             from .audioserver.manager import start_audio_process
 
             process_manager.start_submodule_processes(
@@ -263,7 +263,7 @@ def normal_or_p_d_start(args):
                 ],
             )
 
-        else:
+        elif not args.disable_extra_process_for_multimodal:
             process_manager.start_submodule_processes(
                 start_funcs=[
                     start_visual_process,
diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py
@@ -81,10 +81,12 @@ def __init__(
                 )
 
         self.enable_multimodal = enable_multimodal
+        self.disable_extra_process_for_multimodal = args.disable_extra_process_for_multimodal
         if self.enable_multimodal:
             self.cache_client = rpyc.connect("localhost", cache_port)
-            self.send_to_visual = context.socket(zmq.PUSH)
-            self.send_to_visual.connect(f"{args.zmq_mode}127.0.0.1:{visual_port}")
+            if not self.disable_extra_process_for_multimodal:
+                self.send_to_visual = context.socket(zmq.PUSH)
+                self.send_to_visual.connect(f"{args.zmq_mode}127.0.0.1:{visual_port}")
 
         self.shm_req_manager = ShmReqManager()
 
@@ -449,7 +451,7 @@ async def transfer_to_next_module(
     ):
 
         if self.pd_mode == NodeRole.P:
-            if self.enable_multimodal:
+            if self.enable_multimodal and not self.disable_extra_process_for_multimodal:
                 self.send_to_visual.send_pyobj(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,
@@ -470,7 +472,7 @@ async def transfer_to_next_module(
             return
 
         if self.pd_mode == NodeRole.NORMAL:
-            if self.enable_multimodal:
+            if self.enable_multimodal and not self.disable_extra_process_for_multimodal:
                 self.send_to_visual.send_pyobj(
                     group_req_objs.to_group_req_index(),
                     protocol=pickle.HIGHEST_PROTOCOL,
diff --git a/lightllm/server/multimodal_params.py b/lightllm/server/multimodal_params.py
