1210

Author: sangchengmeng

lightllm/models/qwen3_vl/layer_infer/pre_layer_infer.py

@@ -1,8 +1,94 @@
+import torch
+import torch.distributed as dist
+
+from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+from lightllm.server.embed_cache.utils import (
+    bytes2tensor,
+    read_shm,
+    get_shm_name_embed,
+)
+from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb
+from lightllm.distributed.communication_op import all_reduce
+
 from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer
 
 
 class Qwen3VLMultimodalPreLayerInfer(LlamaMultimodalPreLayerInfer):
     def __init__(self, network_config, mode):
         super().__init__(network_config, mode)
-        self.use_deepstack = True
         return
+
+    def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
+
+        img_weight = []
+        img_start_token_ids = []
+        img_token_lens = []
+        img_start_loc = 0
+        img_start_locs = []
+
+        device = layer_weight.wte_weight_.device
+        dtype = layer_weight.wte_weight_.dtype
+        hidden_size = layer_weight.wte_weight_.shape[1]
+
+        infer_state.mark_multimodal_objs_for_prefill(input_ids=input_ids)
+
+        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 or img["_prefill_"] is False:
+                    continue
+                pos = (input_ids == img["token_id"]).nonzero(as_tuple=True)
+                if pos[0].numel() == 0:
+                    continue
+                # pull the img_embeds by uid from shm
+                all_img_embed_df = bytes2tensor(read_shm(get_shm_name_embed(img["uuid"])))
+                per_image_deepstack = []
+
+                deepstack_layer_num = all_img_embed_df.shape[0] // img["token_num"] - 1
+                img_weight.append(all_img_embed_df[: img["token_num"]].cuda())
+
+                for layer in range(deepstack_layer_num):
+                    start = img["token_num"] * (layer + 1)
+                    end = img["token_num"] * (layer + 2)
+                    per_image_deepstack.append(all_img_embed_df[start:end])
+
+                infer_state.deepstack_features.append(per_image_deepstack)
+                img_insert_locs = int(pos[0][0])
+                infer_state.img_first_token_locs.append(img_insert_locs)
+                infer_state.img_last_token_locs.append(img_insert_locs + img["token_num"])
+                img_start_token_ids.append(img["token_id"])
+                img_token_lens.append(img["token_num"])
+                img_start_locs.append(img_start_loc)
+                img_start_loc += img["token_num"]
+        out = torch.zeros((len(input_ids), hidden_size), dtype=dtype, device=device)
+
+        if len(img_weight) > 0:
+            img_weight = torch.cat(img_weight, dim=0).to(device=device, dtype=dtype)
+        else:
+            img_weight = torch.empty((0, hidden_size), device=device, dtype=dtype)
+        assert img_weight.shape[1] == hidden_size, (
+            f"Dimension mismatch: text weight dimension is {hidden_size}, "
+            f"but image weight dimension is {img_weight.shape[1]}"
+        )
+        # each tp will fill the img embeds, should divide by world_size
+        img_weight = img_weight / self.tp_world_size_
+        img_start_token_ids = torch.Tensor(img_start_token_ids).to(device=device, dtype=torch.long)
+        img_token_lens = torch.Tensor(img_token_lens).to(device=device, dtype=torch.long)
+        img_start_locs = torch.Tensor(img_start_locs).to(device=device, dtype=torch.long)
+
+        multimodal_emb(
+            out,
+            input_ids,
+            layer_weight.wte_weight_,
+            img_weight,
+            img_token_lens,
+            img_start_token_ids,
+            img_start_locs,
+            self.vob_start_id_,
+            self.vob_end_id_,
+        )
+        if self.tp_world_size_ > 1:
+            all_reduce(out, group=infer_state.dist_group, op=dist.ReduceOp.SUM, async_op=False)
+        return out

lightllm/models/qwen3_vl/layer_infer/transformer_layer_infer.py

@@ -55,10 +55,15 @@ def context_forward(self, input_embdings, infer_state: Qwen3VLInferStateInfo, la
                     deepstack_features_cur_layer = deepstack_features[self.layer_num_].to(
                         device=input_embdings.device, non_blocking=True
                     )
+                    print(
+                        f"self.layer_num_ is {self.layer_num_}, i is{i} ,"
+                        f"deepstack_features_cur_layer is {deepstack_features_cur_layer}"
+                    )
                     input_embdings[
                         start:end,
                     ].add_(deepstack_features_cur_layer)
-            infer_state.img_first_token_locs = []
-            infer_state.img_last_token_locs = []
-            infer_state.deepstack_features = []
+            if self.layer_num_ == len(deepstack_features):
+                infer_state.img_first_token_locs = []
+                infer_state.img_last_token_locs = []
+                infer_state.deepstack_features = []
         return input_embdings

lightllm/models/qwen3_vl/model.py

@@ -62,6 +62,7 @@ def get_image_token_length(self, img: ImageItem):
         )
         grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
         token_num = (grid_h * grid_w) // (self.merge_size ** 2)
+        print(f"token_num is {token_num}")
         return token_num
 
     def get_audio_token_length(self, audio: AudioItem):

lightllm/models/qwen3_vl/qwen3_visual.py

@@ -15,6 +15,7 @@
 
 import os
 import json
+import time
 from PIL import Image
 from io import BytesIO
 from typing import List
@@ -67,7 +68,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states.view(
             -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
         )
+        # num_patches = hidden_states.shape[0]
+        # print(f"num_patches is {num_patches}")
+        # torch.cuda.synchronize()
+        # time0 = time.perf_counter()
         hidden_states = self.proj(hidden_states).view(-1, self.embed_dim)
+        # torch.cuda.synchronize()
+        # print(f"patch embed time is {time.perf_counter()-time0}")
         return hidden_states
 
 
@@ -194,6 +201,39 @@ def _init_datatype(self):
             raise ValueError(f"Unsupport datatype {self.data_type}!")
         return
 
+    def concat_img_embed_and_deepstack_features(self, image_embed, deepstack_feature_lists, valid_ids):
+        # input: image_embed: [img_embed1, img_embed2, img_embed3]
+        #        deepstack_feature_lists:[df1-1, df1-2, df1-3,
+        #                                 df2-1, df2-2, df2-3,
+        #                                 df3-1, df3-2, df3-3]
+        #        valid_ids:[[start_1, end_1], [start_2, end_2], [start_3, end_3]]
+        #
+        # return: all_img_embeds_ds: [img_embed1, df1-1, df1-2, df1-3,
+        #                             img_embed2, df2-1, df2-2, df2-3,
+        #                             img_embed3, df3-1, df3-2, df3-3]
+        #         valid_ids:[[start_1, end_1], [start_2, end_2], [start_3, end_3]] # image_embed的start和end
+        all_chunks = []
+        new_valid_ids = []
+
+        row_offset = 0
+
+        for start, end in valid_ids:
+            hs_i = image_embed[start:end]
+            ds_i_list = [feat[start:end] for feat in deepstack_feature_lists]
+
+            combined_i = torch.cat([hs_i, *ds_i_list], dim=0)
+
+            new_start = row_offset
+            new_end = row_offset + combined_i.size(0)
+            new_valid_ids.append([new_start, new_end])
+
+            all_chunks.append(combined_i)
+
+            row_offset += new_end
+
+        all_img_embeds_ds = torch.cat(all_chunks, dim=0)
+        return all_img_embeds_ds, new_valid_ids
+
     def load_model(self, weight_dir):
 
         processor_config_path = os.path.join(weight_dir, "preprocessor_config.json")
@@ -320,21 +360,17 @@ def fast_pos_embed_interpolate(self, grid_thw):
 
     def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
         hidden_states = self.patch_embed(hidden_states)
-
         pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
         hidden_states = hidden_states + pos_embeds
-
         rotary_cos, rotary_sin = self.rot_pos_emb(grid_thw)
         rotary_cos = rotary_cos.to("cuda", non_blocking=True)
         rotary_sin = rotary_sin.to("cuda", non_blocking=True)
-
         cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
             dim=0,
             dtype=torch.int32,
         )
         cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).to("cuda", non_blocking=True)
         max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
-
         deepstack_feature_lists = []
         for layer_num, blk in enumerate(self.blocks):
             hidden_states = blk(
@@ -349,6 +385,7 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs)
                     hidden_states
                 )
                 deepstack_feature_lists.append(deepstack_feature)
+                # print(f"ds time is {time.perf_counter()-time0}")
 
         hidden_states = self.merger(hidden_states)
 
@@ -391,7 +428,9 @@ def encode(self, images: List[ImageItem]):
 
         pixel_values = imgs.to("cuda", dtype=self.data_type, non_blocking=True)
         image_grid_thw = grid_thw.to("cuda", non_blocking=True)
+        img_embeds, deepstack_feature_lists = self.forward(pixel_values, grid_thw=image_grid_thw)
+        all_img_embeds_df, valid_ids = self.concat_img_embed_and_deepstack_features(
+            img_embeds, deepstack_feature_lists, valid_ids
+        )
 
-        all_img_embeds, deepstack_feature_lists = self.forward(pixel_values, grid_thw=image_grid_thw)
-
-        return all_img_embeds, uuids, valid_ids, deepstack_feature_lists
+        return all_img_embeds_df, uuids, valid_ids

lightllm/models/qwen3_vl_moe/layer_infer/transformer_layer_infer.py

@@ -87,7 +87,8 @@ def context_forward(self, input_embdings, infer_state: Qwen3VLInferStateInfo, la
                     input_embdings[
                         start:end,
                     ].add_(deepstack_features_cur_layer)
-            infer_state.img_first_token_locs = []
-            infer_state.img_last_token_locs = []
-            infer_state.deepstack_features = []
+            if self.layer_num_ == len(deepstack_features):
+                infer_state.img_first_token_locs = []
+                infer_state.img_last_token_locs = []
+                infer_state.deepstack_features = []
         return input_embdings

lightllm/models/qwen_vl/layer_infer/pre_layer_infer.py

@@ -6,13 +6,7 @@
 
 from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer
 from lightllm.utils.infer_utils import mark_cost_time
-from lightllm.server.embed_cache.utils import (
-    bytes2tensor,
-    read_shm,
-    get_shm_name_embed,
-    get_shm_name_deepstack,
-    bytes2list,
-)
+from lightllm.server.embed_cache.utils import bytes2tensor, read_shm, get_shm_name_embed
 from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb
 from lightllm.distributed.communication_op import all_reduce
 
@@ -35,7 +29,6 @@
 class LlamaMultimodalPreLayerInfer(LlamaPreLayerInfer):
     def __init__(self, network_config, mode):
         super().__init__(network_config, mode)
-        self.use_deepstack = False
         return
 
     def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight):
@@ -57,18 +50,9 @@ def context_forward(self, input_ids, infer_state: LlamaInferStateInfo, layer_wei
                 # skip the same image
                 if img["token_id"] in img_start_token_ids or img["_prefill_"] is False:
                     continue
-                pos = (input_ids == img["token_id"]).nonzero(as_tuple=True)
-                if pos[0].numel() == 0:
-                    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.use_deepstack:
-                    deepstack_features = read_shm(get_shm_name_deepstack(img["uuid"]))
-                    infer_state.deepstack_features.append(bytes2list(deepstack_features))
-                    img_insert_locs = int(pos[0][0])
-                    infer_state.img_first_token_locs.append(img_insert_locs)
-                    infer_state.img_last_token_locs.append(img_insert_locs + img["token_num"])
                 img_start_token_ids.append(img["token_id"])
                 img_token_lens.append(img["token_num"])
                 img_start_locs.append(img_start_loc)

lightllm/server/embed_cache/utils.py

@@ -21,52 +21,11 @@ def tensor2bytes(t: torch.Tensor):
     return buf.read()
 
 
-def list2bytes(tensors: List[torch.Tensor]) -> bytes:
-    # 逐个张量做 detach().cpu() 和复制
-    safe_list = []
-    for t in tensors:
-        if t is None:
-            safe_list.append(None)
-            continue
-        t = t.detach().cpu()
-        if not t.is_contiguous():
-            t = t.contiguous()
-        dest = torch.empty_like(t)
-        dest.copy_(t)
-        safe_list.append(dest)
-    buf = BytesIO()
-    torch.save(safe_list, buf, _use_new_zipfile_serialization=False, pickle_protocol=4)
-    buf.seek(0)
-    return buf.read()
-
-
 def bytes2tensor(b):
     # return torch.from_numpy(np.frombuffer(b, dtype=np.float16)).cuda()
     return torch.load(BytesIO(b), weights_only=False)
 
 
-def bytes2list(b: bytes, device: Optional[torch.device] = None, non_blocking: bool = False) -> List[torch.Tensor]:
-    obj = torch.load(BytesIO(b), map_location="cpu", weights_only=False)
-
-    if isinstance(obj, tuple):
-        obj = list(obj)
-    if not isinstance(obj, list):
-        raise TypeError(f"Loaded object is {type(obj)}, expected list or tuple.")
-
-    if device is None:
-        return obj
-
-    out: List[torch.Tensor] = []
-    for x in obj:
-        if x is None:
-            out.append(None)
-        elif isinstance(x, torch.Tensor):
-            out.append(x.to(device, non_blocking=non_blocking))
-        else:
-            raise TypeError(f"List element is {type(x)}, expected Tensor or None.")
-    return out
-
-
 def create_shm(name, data):
     try:
         data_size = len(data)
@@ -95,7 +54,3 @@ def get_shm_name_data(uid):
 
 def get_shm_name_embed(uid):
     return str(uid) + "-embed"
-
-
-def get_shm_name_deepstack(uid):
-    return str(uid) + "-deepstack"

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -21,12 +21,8 @@
 from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel
 from lightllm.server.embed_cache.utils import (
     tensor2bytes,
-    read_shm,
     create_shm,
-    get_shm_name_data,
     get_shm_name_embed,
-    get_shm_name_deepstack,
-    list2bytes,
 )
 from lightllm.utils.infer_utils import set_random_seed
 from lightllm.utils.infer_utils import calculate_time, mark_start, mark_end
@@ -111,8 +107,7 @@ def forward(self, images: List[ImageItem]):
     # @calculate_time(show=False, min_cost_ms=300)
     def exposed_encode(self, images: List[ImageItem]):
         images = obtain(images)
-        all_img_embeds, uuids, valid_ids, *deepstack_features = self.forward(images)
-        deepstack_feature_lists = deepstack_features[0] if deepstack_features else None
+        all_img_embeds, uuids, valid_ids = self.forward(images)
         all_img_embeds = all_img_embeds.to(torch.device("cpu"))
 
         if self.tp_rank_id == 0:
@@ -125,10 +120,6 @@ def exposed_encode(self, images: List[ImageItem]):
                 start, end = valid_ids[i]
                 cur_embed_bytes = tensor2bytes(all_img_embeds[start:end])
                 create_shm(get_shm_name_embed(uid), cur_embed_bytes)
-                if deepstack_feature_lists is not None:
-                    per_image_deepstack = [feat[start:end] for feat in deepstack_feature_lists]
-                    deepstack_features_bytes = list2bytes(per_image_deepstack)
-                    create_shm(get_shm_name_deepstack(uid), deepstack_features_bytes)
                 ids_to_set.append(uid)
             if ids_to_set:
                 self.cache_client.root.set_items_embed(ids_to_set)