sglang 0.4.4.post3__py3-none-any.whl → 0.4.5__py3-none-any.whl

sglang/srt/models/deepseek_vl2.py

@@ -11,7 +11,11 @@ from sglang.srt.configs.deepseekvl2 import (
 )
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternImageTokens,
+    general_mm_embed_routine,
+)
+from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
@@ -150,7 +154,6 @@ class DeepseekVL2MlpProjector(nn.Module):
         return x
 
 
-# todo
 class DeepseekVL2ForCausalLM(nn.Module):
 
     def __init__(
@@ -215,32 +218,15 @@ class DeepseekVL2ForCausalLM(nn.Module):
         forward_batch: ForwardBatch,
         **kwargs: object,
     ):
-        input_embeds = self.language_model.model.embed_tokens(input_ids)
-        if (
-            forward_batch.forward_mode.is_extend()
-            and forward_batch.contains_image_inputs()
-        ):
-            extend_start_loc_cpu = forward_batch.extend_start_loc.cpu().numpy()
-            extend_seq_lens_cpu = forward_batch.extend_seq_lens.cpu().numpy()
-            for idx, image in enumerate(forward_batch.mm_inputs):
-                if image is None:
-                    continue
-                start_idx = extend_start_loc_cpu[idx]
-                end_idx = start_idx + extend_seq_lens_cpu[idx]
-                images_emb_mask = image.images_emb_mask.to(device="cuda")
-                image_features = self.get_image_feature(image)
-                input_embeds[start_idx:end_idx] = input_embeds[
-                    start_idx:end_idx
-                ].masked_scatter(images_emb_mask.unsqueeze(-1), image_features)
-
-        outputs = self.language_model.forward(
+        hs = general_mm_embed_routine(
             input_ids=input_ids,
             positions=positions,
             forward_batch=forward_batch,
-            input_embeds=input_embeds,
+            image_data_embedding_func=self.get_image_feature,
+            language_model=self.language_model,
         )
 
-        return outputs
+        return hs
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
@@ -263,94 +249,109 @@ class DeepseekVL2ForCausalLM(nn.Module):
                 weights_loader(param, loaded_weight)
 
     def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs):
-        return input_ids
-
-    def get_image_feature(self, image_input: MultimodalInputs):
-        pixel_values = image_input.pixel_values.type(
-            next(self.vision.parameters()).dtype
-        ).to(device=next(self.vision.parameters()).device)
-        image_feature = self.vision.forward_features(pixel_values)
-        images_embeds = self.projector(image_feature)
-        _, hw, n_dim = images_embeds.shape
-        h = w = int(hw**0.5)
-        tile_index = 0
+        helper = MultiModalityDataPaddingPatternImageTokens(
+            image_token_id=image_inputs.im_token_id
+        )
+        return helper.pad_input_tokens(input_ids, image_inputs)
+
+    def get_image_feature(self, items: List[MultimodalDataItem]):
+
+        images_spatial_crop = torch.cat(
+            [item.image_spatial_crop for item in items], dim=0
+        )
+
+        assert images_spatial_crop.dim() == 3
+
+        # TODO: can it be batched ?
         images_in_this_batch = []
-        images_spatial_crop = image_input.image_spatial_crop
-        for jdx in range(images_spatial_crop.shape[1]):
-            num_width_tiles, num_height_tiles = images_spatial_crop[0, jdx]
-            if num_width_tiles == 0 or num_height_tiles == 0:
-                break
-            num_tiles_in_image = num_width_tiles * num_height_tiles
-
-            # [hw, D]
-            global_features = images_embeds[tile_index]
-
-            # [num_height_tiles * num_width_tiles, hw, D]
-            local_features = images_embeds[
-                tile_index + 1 : tile_index + 1 + num_tiles_in_image
-            ]
-            tile_index += num_tiles_in_image + 1
-
-            # format global and local features
-            # ----------------- global view add newline -----------------
-            # [hw, D] -> [h, w, D]
-            global_features = global_features.view(h, w, n_dim)
-
-            # [D]     -> [h, 1, D]
-            new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h)
-
-            # cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D]
-            global_features = torch.cat([global_features, new_lines_in_global], dim=1)
-
-            # [h, w + 1, D] -> [h * (w + 1), D]
-            global_features = global_features.view(-1, n_dim)
-
-            # ----------------- local view add newline -----------------
-            # [num_height_tiles * num_width_tiles, h * w, D] ->
-            # [num_height_tiles * h, num_width_tiles * w, D]
-            local_features = rearrange(
-                local_features,
-                "(th tw) (h w) d -> (th h) (tw w) d",
-                th=num_height_tiles,
-                tw=num_width_tiles,
-                h=h,
-                w=w,
+        for item in items:
+            assert item.pixel_values.dim() == 4
+            image_feature = self.vision.forward_features(
+                item.pixel_values.type(next(self.vision.parameters()).dtype).to(
+                    device=next(self.vision.parameters()).device
+                )
             )
+            images_embeds = self.projector(image_feature)
+            _, hw, n_dim = images_embeds.shape
+            h = w = int(hw**0.5)
+            tile_index = 0
+            for jdx in range(item.image_spatial_crop.shape[1]):
+                num_width_tiles, num_height_tiles = item.image_spatial_crop[0, jdx]
+                if num_width_tiles == 0 or num_height_tiles == 0:
+                    break
+                num_tiles_in_image = num_width_tiles * num_height_tiles
+
+                # [hw, D]
+                global_features = images_embeds[tile_index]
+
+                # [num_height_tiles * num_width_tiles, hw, D]
+                local_features = images_embeds[
+                    tile_index + 1 : tile_index + 1 + num_tiles_in_image
+                ]
+                tile_index += num_tiles_in_image + 1
 
-            # [D] -> [num_height_tiles * h, 1, D]
-            new_lines_in_local = repeat(
-                self.image_newline,
-                "d -> (th h) 1 d",
-                th=num_height_tiles,
-                h=h,
-            )
+                # format global and local features
+                # ----------------- global view add newline -----------------
+                # [hw, D] -> [h, w, D]
+                global_features = global_features.view(h, w, n_dim)
+
+                # [D]     -> [h, 1, D]
+                new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h)
 
-            # [num_height_tiles * h, num_width_tiles * w + 1, D]
-            local_features = torch.cat([local_features, new_lines_in_local], dim=1)
-
-            # [num_height_tiles * h, num_width_tiles * w + 1, D]
-            #   --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D]
-            local_features = local_features.view(-1, n_dim)
-
-            # merge global and local tiles
-            if self.global_view_pos == "head":
-                global_local_features = torch.cat(
-                    [
-                        global_features,
-                        self.view_seperator[None, :],
-                        local_features,
-                    ]
+                # cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D]
+                global_features = torch.cat(
+                    [global_features, new_lines_in_global], dim=1
                 )
-            else:
-                global_local_features = torch.cat(
-                    [
-                        local_features,
-                        self.view_seperator[None, :],
-                        global_features,
-                    ]
+
+                # [h, w + 1, D] -> [h * (w + 1), D]
+                global_features = global_features.view(-1, n_dim)
+
+                # ----------------- local view add newline -----------------
+                # [num_height_tiles * num_width_tiles, h * w, D] ->
+                # [num_height_tiles * h, num_width_tiles * w, D]
+                local_features = rearrange(
+                    local_features,
+                    "(th tw) (h w) d -> (th h) (tw w) d",
+                    th=num_height_tiles,
+                    tw=num_width_tiles,
+                    h=h,
+                    w=w,
                 )
 
-            images_in_this_batch.append(global_local_features)
+                # [D] -> [num_height_tiles * h, 1, D]
+                new_lines_in_local = repeat(
+                    self.image_newline,
+                    "d -> (th h) 1 d",
+                    th=num_height_tiles,
+                    h=h,
+                )
+
+                # [num_height_tiles * h, num_width_tiles * w + 1, D]
+                local_features = torch.cat([local_features, new_lines_in_local], dim=1)
+
+                # [num_height_tiles * h, num_width_tiles * w + 1, D]
+                #   --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D]
+                local_features = local_features.view(-1, n_dim)
+
+                # merge global and local tiles
+                if self.global_view_pos == "head":
+                    global_local_features = torch.cat(
+                        [
+                            global_features,
+                            self.view_seperator[None, :],
+                            local_features,
+                        ]
+                    )
+                else:
+                    global_local_features = torch.cat(
+                        [
+                            local_features,
+                            self.view_seperator[None, :],
+                            global_features,
+                        ]
+                    )
+
+                images_in_this_batch.append(global_local_features)
 
         return torch.cat(images_in_this_batch, dim=0)
 

sglang/srt/models/gemma3_mm.py

@@ -21,14 +21,7 @@ from typing import Dict, Iterable, List, Optional, Set, Tuple, TypedDict
 
 import torch
 from torch import nn
-from transformers import (
-    AutoModel,
-    BatchFeature,
-    Gemma3Config,
-    Gemma3Processor,
-    PreTrainedModel,
-)
-from transformers.models.gemma3.processing_gemma3 import Gemma3ProcessorKwargs
+from transformers import AutoModel, Gemma3Config, PreTrainedModel
 
 from sglang.srt.hf_transformers_utils import get_processor
 from sglang.srt.layers.layernorm import Gemma3RMSNorm
@@ -38,7 +31,11 @@ from sglang.srt.managers.mm_utils import (
     MultiModalityDataPaddingPatternTokenPairs,
     general_mm_embed_routine,
 )
-from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.managers.schedule_batch import (
+    MultimodalDataItem,
+    MultimodalInputs,
+    flatten_nested_list,
+)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
@@ -274,17 +271,16 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
         """
         return self.language_model.get_attention_sliding_window_size()
 
-    def get_image_feature(self, image_input: MultimodalInputs):
+    def get_image_feature(self, items: List[MultimodalDataItem]):
         """
         Projects the last hidden state from the vision model into language model space.
 
-        Args:
-            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)
-               The tensors corresponding to the input images.
         Returns:
             image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).
         """
-        pixel_values = image_input.pixel_values
+        pixel_values = torch.stack(
+            flatten_nested_list([item.pixel_values for item in items]), dim=0
+        )
         pixel_values = pixel_values.to("cuda")
         pixel_values = pixel_values.to(dtype=self.language_model.dtype())
 
@@ -292,61 +288,6 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
         image_features = self.multi_modal_projector(vision_outputs)
         return image_features
 
-    def embed_mm_inputs(
-        self,
-        input_ids: torch.Tensor,
-        forward_batch: ForwardBatch,
-        image_input: MultimodalInputs,
-    ) -> torch.Tensor:
-        if input_ids is None:
-            raise ValueError("Unimplemented")
-        # boolean-masking image tokens
-        special_image_mask = torch.isin(
-            input_ids,
-            torch.tensor(image_input.pad_values, device=input_ids.device),
-        ).unsqueeze(-1)
-        num_image_tokens_in_input_ids = special_image_mask.sum()
-
-        inputs_embeds = None
-        if num_image_tokens_in_input_ids == 0:
-            inputs_embeds = self.get_input_embeddings()(input_ids)
-            return inputs_embeds
-        else:
-            # print(f"image tokens from input_ids: {inputs_embeds[special_image_mask].numel()}")
-            image_features = self.get_image_feature(image_input.pixel_values)
-
-            # print(f"image tokens from image embeddings: {image_features.numel()}")
-            num_image_tokens_in_embedding = (
-                image_features.shape[0] * image_features.shape[1]
-            )
-
-            if num_image_tokens_in_input_ids != num_image_tokens_in_embedding:
-                num_image = num_image_tokens_in_input_ids // image_features.shape[1]
-                image_features = image_features[:num_image, :]
-                logger.warning(
-                    f"Number of images does not match number of special image tokens in the input text. "
-                    f"Got {num_image_tokens_in_input_ids} image tokens in the text but {num_image_tokens_in_embedding} "
-                    "tokens from image embeddings."
-                )
-
-            # Important: clamp after extracting original image boundaries
-            input_ids.clamp_(min=0, max=self.vocab_size - 1)
-
-            inputs_embeds = self.get_input_embeddings()(input_ids)
-
-            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(
-                inputs_embeds.device
-            )
-
-            image_features = image_features.to(
-                inputs_embeds.device, inputs_embeds.dtype
-            )
-            inputs_embeds = inputs_embeds.masked_scatter(
-                special_image_mask, image_features
-            )
-
-        return inputs_embeds
-
     @torch.no_grad()
     def forward(
         self,
@@ -405,22 +346,15 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
         else:
             llm_input_ids = input_ids
 
-        inputs_embeds = general_mm_embed_routine(
+        hs = general_mm_embed_routine(
             input_ids=llm_input_ids,
             forward_batch=forward_batch,
-            embed_tokens=self.get_input_embeddings(),
-            mm_data_embedding_func=self.get_image_feature,
-        )
-
-        outputs = self.language_model(
-            input_ids=None,
+            language_model=self.language_model,
+            image_data_embedding_func=self.get_image_feature,
             positions=positions,
-            forward_batch=forward_batch,
-            input_embeds=inputs_embeds,
-            **kwargs,
         )
 
-        return outputs
+        return hs
 
     def tie_weights(self):
         return self.language_model.tie_weights()

sglang/srt/models/llama.py

@@ -17,7 +17,7 @@
 """Inference-only LLaMA model compatible with HuggingFace weights."""
 
 import logging
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -63,6 +63,7 @@ class LlamaMLP(nn.Module):
         hidden_act: str,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        reduce_results: bool = True,
     ) -> None:
         super().__init__()
         self.gate_up_proj = MergedColumnParallelLinear(
@@ -78,6 +79,7 @@ class LlamaMLP(nn.Module):
             bias=False,
             quant_config=quant_config,
             prefix=add_prefix("down_proj", prefix),
+            reduce_results=reduce_results,
         )
         if hidden_act != "silu":
             raise ValueError(
@@ -281,7 +283,7 @@ class LlamaModel(nn.Module):
         self.layers = make_layers(
             config.num_hidden_layers,
             lambda idx, prefix: LlamaDecoderLayer(
-                config=config, quant_config=quant_config, layer_id=idx, prefix=prefix
+                config=config, layer_id=idx, quant_config=quant_config, prefix=prefix
             ),
             prefix="model.layers",
         )
@@ -375,9 +377,7 @@ class LlamaForCausalLM(nn.Module):
         super().__init__()
         self.config = config
         self.quant_config = quant_config
-        self.model = LlamaModel(
-            config, quant_config=quant_config, prefix=add_prefix("model", prefix)
-        )
+        self.model = self._init_model(config, quant_config, add_prefix("model", prefix))
         # Llama 3.2 1B Instruct set tie_word_embeddings to True
         # Llama 3.1 8B Instruct set tie_word_embeddings to False
         if self.config.tie_word_embeddings:
@@ -402,6 +402,14 @@ class LlamaForCausalLM(nn.Module):
 
         self.capture_aux_hidden_states = False
 
+    def _init_model(
+        self,
+        config: LlamaConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
+        return LlamaModel(config, quant_config=quant_config, prefix=prefix)
+
     @torch.no_grad()
     def forward(
         self,
@@ -428,6 +436,9 @@ class LlamaForCausalLM(nn.Module):
         else:
             return self.pooler(hidden_states, forward_batch)
 
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
     def get_hidden_dim(self, module_name):
         # return input_dim, output_dim
         if module_name in ["q_proj", "o_proj", "qkv_proj"]: