diff --git a/docs/supported_models/multimodal_language_models.md b/docs/supported_models/multimodal_language_models.md index 792c67e30..1677bb574 100644 --- a/docs/supported_models/multimodal_language_models.md +++ b/docs/supported_models/multimodal_language_models.md @@ -46,6 +46,7 @@ in the GitHub search bar. | **DotsVLM-OCR** | `rednote-hilab/dots.ocr` | Specialized OCR variant of DotsVLM optimized for optical character recognition tasks with enhanced text extraction and document understanding capabilities. | Don't use `--trust-remote-code` | | **NVILA** (8B, 15B, Lite-2B, Lite-8B, Lite-15B) | `Efficient-Large-Model/NVILA-8B` | `chatml` | NVILA explores the full stack efficiency of multi-modal design, achieving cheaper training, faster deployment and better performance. | | **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | NVIDIA Nemotron Nano v2 VL enables multi-image reasoning and video understanding, along with strong document intelligence, visual Q&A and summarization capabilities. It builds on Nemotron Nano V2, a hybrid Mamba-Transformer LLM, in order to achieve higher inference throughput in long document and video scenarios. | Use `--trust-remote-code`. You may need to adjust `--max-mamba-cache-size` [default is 512] to fit memory constraints. | +| **Ernie4.5-VL** | `baidu/ERNIE-4.5-VL-28B-A3B-PT` | Baidu's vision-language models(28B,424B). Support image and video comprehension, and also support thinking. | | | **JetVLM** | | JetVLM is an vision-language model designed for high-performance multimodal understanding and generation tasks built upon Jet-Nemotron. | Coming soon | ## Video Input Support diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 1780f844e..f9d5a831d 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1123,6 +1123,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal multimodal_model_archs = [ "CLIPModel", "DeepseekVL2ForCausalLM", + "Ernie4_5_VLMoeForConditionalGeneration", "Gemma3ForConditionalGeneration", "Gemma3nForConditionalGeneration", "Glm4vForConditionalGeneration", diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 064d55743..b38c980f1 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -2285,6 +2285,177 @@ class MRotaryEmbedding(RotaryEmbedding): return position_ids, mrope_position_deltas + @staticmethod + def get_rope_index_ernie45( + input_ids: torch.Tensor, + hf_config: Any, + image_grid_thw: Union[list[list[int]], torch.Tensor], + video_grid_thw: Union[list[list[int]], torch.Tensor], + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Get mrope input positions and delta value for Ernie VL.""" + + image_token_id = hf_config.im_patch_id + video_start_token_id = hf_config.video_start_token_id + video_end_token_id = hf_config.video_end_token_id + spatial_conv_size = hf_config.spatial_conv_size + temporal_conv_size = hf_config.temporal_conv_size + + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + for i, input_ids in enumerate(total_input_ids): + input_tokens = input_ids.tolist() + + input_token_type = [] + video_check_flg = False + for token in input_tokens: + if token == video_start_token_id: + video_check_flg = True + elif token == video_end_token_id: + video_check_flg = False + + if token == image_token_id and not video_check_flg: + input_token_type.append("image") + elif token == image_token_id and video_check_flg: + input_token_type.append("video") + else: + input_token_type.append("text") + + input_type_group = [] + for key, group in itertools.groupby( + enumerate(input_token_type), lambda x: x[1] + ): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + llm_pos_ids_list = [] + video_frame_num = 1 + for modality_type, start_idx, end_idx in input_type_group: + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + + if modality_type == "image": + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_conv_size, + w.item() // spatial_conv_size, + ) + + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + st_idx + ) + + image_index += 1 + video_frame_num = 1 + + elif modality_type == "video": + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item() // temporal_conv_size, + h.item() // spatial_conv_size, + w.item() // spatial_conv_size, + ) + + for t_idx in range(llm_grid_t): + t_index = ( + torch.tensor(t_idx) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(1, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(1, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + st_idx + ) + + video_index += 1 + video_frame_num += 1 + + else: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + video_frame_num = 1 + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, :] = llm_positions.to(position_ids.device) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + s = input_ids.shape[1] + position_ids = torch.arange(s) + position_ids = ( + position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + 1 - s + return position_ids, mrope_position_deltas + # For qwen3-omni @staticmethod def _get_feat_extract_output_lengths(input_lengths): @@ -2324,6 +2495,91 @@ class MRotaryEmbedding(RotaryEmbedding): return llm_pos_ids +class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding): + """3D rotary positional embedding. [h w h w h w h w... t t t...]""" + + def forward_native( # type: ignore[override] + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert positions.ndim == 1 or positions.ndim == 2 + assert key is not None + + num_tokens = positions.shape[-1] + cos_sin = self.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if positions.ndim == 2: + assert self.mrope_section + + section_h = self.mrope_section[0] # 22 + section_w = self.mrope_section[1] # 22 + section_t = self.mrope_section[2] # 20 + assert section_h == section_w + # Split according to [h w h w h w h w... t t t...] + section_cos_t = cos[..., -section_t:] + section_cos_h = cos[..., : section_h + section_w : 2] + section_cos_w = cos[..., 1 : section_h + section_w : 2] + + cos_t, cos_h, cos_w = section_cos_t[0], section_cos_h[1], section_cos_w[2] + cos_hw = torch.stack([cos_h, cos_w], dim=-1).reshape( + cos_h.shape[:-1] + (cos_h.shape[-1] * 2,) + ) + cos = torch.cat([cos_hw, cos_t], dim=-1) + + section_sin_t = sin[..., -section_t:] + section_sin_h = sin[..., : section_h + section_w : 2] + section_sin_w = sin[..., 1 : section_h + section_w : 2] + + sin_t, sin_h, sin_w = section_sin_t[0], section_sin_h[1], section_sin_w[2] + sin_hw = torch.stack([sin_h, sin_w], dim=-1).reshape( + sin_h.shape[:-1] + (sin_h.shape[-1] * 2,) + ) + sin = torch.cat([sin_hw, sin_t], dim=-1) + + query_shape = query.shape + query = query.view(num_tokens, -1, self.head_size) + query_rot = query[..., : self.rotary_dim] + query_pass = query[..., self.rotary_dim :] + query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + key = key.view(num_tokens, -1, self.head_size) + key_rot = key[..., : self.rotary_dim] + key_pass = key[..., self.rotary_dim :] + key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + return query, key + + def forward_cuda( # type: ignore[override] + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + return self.forward_native(positions, query, key) + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass with optional Triton kernel acceleration. + Args: + positions: + [num_tokens,] (text only) or + [3, num_tokens] (T/H/W positions with multimodal inputs) + query: [num_tokens, num_heads * head_size] + key: [num_tokens, num_kv_heads * head_size] + """ + assert positions.ndim == 1 or positions.ndim == 2 + return self.forward_native(positions, query, key) + + class DualChunkRotaryEmbedding(MultiPlatformOp): """Rotary positional embedding for Dual Chunk Attention.""" diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py new file mode 100644 index 000000000..3fe0fc6a7 --- /dev/null +++ b/python/sglang/srt/models/ernie45_moe_vl.py @@ -0,0 +1,552 @@ +# Copyright 2023-2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +""" Inference-only Ernie4.5 VL model compatible with baidu/ERNIE-4.5-VL-*-PT weights. """ + +import logging +from itertools import islice +from typing import Any, Dict, Optional, Tuple, Union + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import Ernie4_5_VLRotaryEmbedding +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP +from sglang.srt.utils import add_prefix, make_layers + +logger = logging.getLogger(__name__) + + +class Ernie4_5_VLMoeAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + rope_is_neox_style: bool = True, + freq_allocation: int = 20, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + bias: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + # 3D rope + t_rope = freq_allocation + h_rope = (self.head_dim // 2 - freq_allocation) // 2 + w_rope = (self.head_dim // 2 - freq_allocation) // 2 + + self.rotary_emb = Ernie4_5_VLRotaryEmbedding( + head_size=self.head_dim, + rotary_dim=self.head_dim, + max_position_embeddings=max_position_embeddings, + base=rope_theta, + is_neox_style=False, + dtype=torch.get_default_dtype(), + mrope_section=[h_rope, w_rope, t_rope], + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class Ernie4_5_VLMoeMoE(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.layer_id = layer_id + self.tp_size = get_tensor_model_parallel_world_size() + self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", 0) + self.hidden_size = config.hidden_size + + moe_num_experts = config.moe_num_experts + max_moe_num_experts = max(moe_num_experts) + + if self.tp_size > max_moe_num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {moe_num_experts}." + ) + + moe_layer_start_index = config.moe_layer_start_index + text_moe_layer_start_index = moe_layer_start_index[0] + vision_moe_layer_start_index = moe_layer_start_index[1] + moe_layer_end_index = config.moe_layer_end_index + moe_layer_end_index = getattr( + config, + "moe_layer_end_index", + [config.num_hidden_layers - 1, config.num_hidden_layers - 1], + ) + text_moe_layer_end_index = moe_layer_end_index[0] + vision_moe_layer_end_index = moe_layer_end_index[1] + + assert config.moe_num_experts[0] == config.moe_num_experts[1] + self.e_score_correction_bias = nn.Parameter( + torch.empty(2, config.moe_num_experts[0], dtype=torch.float32) + ) + + assert text_moe_layer_start_index <= text_moe_layer_end_index + + if ( + layer_id >= text_moe_layer_start_index + and layer_id <= text_moe_layer_end_index + ): + self.text_experts_gate = ReplicatedLinear( + config.hidden_size, + config.moe_num_experts[0], + bias=False, + params_dtype=torch.float32, + quant_config=quant_config, + prefix=add_prefix("text_experts_gate", prefix), + ) + + self.text_experts_topk = TopK( + top_k=config.moe_k, + renormalize=True, + use_grouped_topk=False, + correction_bias=self.e_score_correction_bias[0], + ) + + self.text_experts = get_moe_impl_class(quant_config)( + num_experts=config.moe_num_experts[0], + top_k=config.moe_k, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size[0], + layer_id=self.layer_id, + quant_config=quant_config, + prefix=add_prefix("text_experts", prefix), + ) + + assert vision_moe_layer_start_index <= vision_moe_layer_end_index + if ( + layer_id >= vision_moe_layer_start_index + and layer_id <= vision_moe_layer_end_index + ): + + self.vision_experts_gate = ReplicatedLinear( + config.hidden_size, + config.moe_num_experts[1], + bias=False, + params_dtype=torch.float32, + quant_config=quant_config, + prefix=add_prefix("vision_experts_gate", prefix), + ) + + self.vision_experts_topk = TopK( + top_k=config.moe_k, + renormalize=True, + use_grouped_topk=False, + correction_bias=self.e_score_correction_bias[1], + ) + + self.vision_experts = get_moe_impl_class(quant_config)( + num_experts=config.moe_num_experts[1], + top_k=config.moe_k, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size[1], + layer_id=self.layer_id, + quant_config=quant_config, + prefix=add_prefix("vision_experts", prefix), + ) + + if self.moe_num_shared_experts > 0: + intermediate_size = ( + config.moe_intermediate_size[0] * config.moe_num_shared_experts + ) + self.shared_experts = Ernie4_5_VLMoeMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + visual_token_mask: torch.Tensor, + **kwargs: object, + ) -> torch.Tensor: + shared_output = ( + self.shared_experts(hidden_states) + if self.moe_num_shared_experts > 0 + else None + ) + + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + + capturing = torch.cuda.is_current_stream_capturing() + + if visual_token_mask is not None and not capturing: + all_visual = visual_token_mask.all() + any_visual = visual_token_mask.any() + else: + # During CUDA Graph capture, all set false + all_visual = False + any_visual = False + + if all_visual: + # vision modal input processing directly + vision_router_logits, _ = self.vision_experts_gate( + hidden_states.to(dtype=torch.float32) + ) + vision_topk_output = self.vision_experts_topk( + hidden_states, vision_router_logits + ) + final_hidden_states = self.vision_experts( + hidden_states=hidden_states, topk_output=vision_topk_output + ) + elif any_visual: + visual_token_mask = visual_token_mask.repeat(1, self.hidden_size).bool() + text_token_mask = ~visual_token_mask + final_hidden_states = torch.zeros_like(hidden_states) + + text_hidden_states = hidden_states[text_token_mask].reshape( + -1, self.hidden_size + ) + vision_hidden_states = hidden_states[visual_token_mask].reshape( + -1, self.hidden_size + ) + + text_router_logits, _ = self.text_experts_gate( + text_hidden_states.to(dtype=torch.float32) + ) + text_topk_output = self.text_experts_topk( + text_hidden_states, text_router_logits + ) + final_hidden_states[text_token_mask] = self.text_experts( + hidden_states=text_hidden_states, topk_output=text_topk_output + ).flatten() + + vision_router_logits, _ = self.vision_experts_gate( + vision_hidden_states.to(dtype=torch.float32) + ) + vision_topk_output = self.vision_experts_topk( + vision_hidden_states, vision_router_logits + ) + final_hidden_states[visual_token_mask] = self.vision_experts( + hidden_states=vision_hidden_states, topk_output=vision_topk_output + ).flatten() + + else: + # text modal input processing directly + text_router_logits, _ = self.text_experts_gate( + hidden_states.to(dtype=torch.float32) + ) + topk_output = self.text_experts_topk(hidden_states, text_router_logits) + final_hidden_states = self.text_experts( + hidden_states=hidden_states, topk_output=topk_output + ) + + if shared_output is not None: + final_hidden_states = final_hidden_states + shared_output + + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(orig_shape) + + +class Ernie4_5_VLMoeDecoderLayer(nn.Module): + """A single transformer layer. + + Transformer layer takes input with size [s, b, h] and returns an + output of the same size. + """ + + def __init__( + self, + config, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + rope_theta = getattr(config, "rope_theta", 500000) + rope_scaling = getattr(config, "rope_scaling", None) + rope_is_neox_style = getattr(config, "rope_is_neox_style", False) + freq_allocation = getattr(config, "freq_allocation", 20) + max_position_embeddings = getattr(config, "max_position_embeddings", 131072) + # Self attention. + self.self_attn = Ernie4_5_VLMoeAttention( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + rope_is_neox_style=rope_is_neox_style, + freq_allocation=freq_allocation, + max_position_embeddings=config.max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + bias=config.use_bias, + ) + + # MoE + moe_layer_start_index = config.moe_layer_start_index + min_moe_layer_start_index = min(moe_layer_start_index) + moe_layer_end_index = getattr( + config, + "moe_layer_end_index", + [config.num_hidden_layers - 1, config.num_hidden_layers - 1], + ) + max_moe_layer_end_index = max(moe_layer_end_index) + assert min_moe_layer_start_index <= max_moe_layer_end_index + moe_num_experts = config.moe_num_experts + max_moe_num_experts = max(moe_num_experts) + moe_layer_interval = getattr(config, "moe_layer_interval", 1) + use_moe = getattr(config, "use_moe", max_moe_num_experts > 0) + # MLP + if ( + use_moe + and ((layer_id + 1) % moe_layer_interval == 0) + and layer_id >= min_moe_layer_start_index + and layer_id <= max_moe_layer_end_index + ): + self.mlp = Ernie4_5_VLMoeMoE( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + self.mlp = Ernie4_5_VLMoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + visual_token_mask: torch.Tensor | None, + **kwargs: object, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + if isinstance(self.mlp, Ernie4_5_VLMoeMoE): + hidden_states = self.mlp(hidden_states, visual_token_mask, **kwargs) + else: + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual + + +# only used as text backbone for ernie4.5 vl +class Ernie4_5_VLMoeModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Ernie4_5_VLMoeDecoderLayer( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def get_input_embeddings(self) -> torch.Tensor: + return self.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + visual_token_mask: torch.Tensor | None = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + visual_token_mask, + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + + if hidden_states.shape[0] != 0: + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + return hidden_states diff --git a/python/sglang/srt/models/ernie45_vl.py b/python/sglang/srt/models/ernie45_vl.py new file mode 100644 index 000000000..9ce3e97ce --- /dev/null +++ b/python/sglang/srt/models/ernie45_vl.py @@ -0,0 +1,845 @@ +# Copyright 2023-2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only Ernie45-VL model compatible with HuggingFace weights.""" +import logging +from functools import lru_cache, partial +from typing import Iterable, List, Optional, Tuple, Type + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers import PretrainedConfig + +from sglang.srt.layers.activation import QuickGELU +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + 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.ernie45_moe_vl import Ernie4_5_VLMoeModel +from sglang.srt.utils import add_prefix +from sglang.srt.utils.hf_transformers_utils import get_processor + +logger = logging.getLogger(__name__) + + +# === Vision Encoder === # + + +class Ernie4_5_VisionMLP(nn.Module): + + def __init__( + self, + in_features: int, + hidden_features: int = None, + act_layer: Type[nn.Module] = QuickGELU, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.fc1 = ColumnParallelLinear( + in_features, + hidden_features, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + self.act = act_layer() + self.fc2 = RowParallelLinear( + hidden_features, + in_features, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x_parallel, _ = self.fc1(x) + x_parallel = self.act(x_parallel) + x, _ = self.fc2(x_parallel) + return x + + +class Ernie4_5_VisionBlock(nn.Module): + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float, + act_layer: Type[nn.Module] = QuickGELU, + norm_layer: Type[nn.Module] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + use_qkv_parallel=True, + flatten_batch=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.mlp = Ernie4_5_VisionMLP( + dim, + mlp_hidden_dim, + act_layer=act_layer, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.norm1(x) + hidden_states = rearrange(hidden_states, "s b ... -> b s ...") + attn = self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + ) + attn = rearrange(attn, "b s ... -> s b ...") + x = x + attn + x = x + self.mlp(self.norm2(x)) + return x + + +class Ernie4_5_VisionPatchEmbed(nn.Module): + + def __init__( + self, + patch_size: int = 14, + in_chans: int = 3, + embed_dim: int = 1280, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.in_channels = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Linear(in_chans * patch_size * patch_size, embed_dim, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.to(target_dtype) + hidden_states = self.proj(hidden_states) + + return hidden_states + + +class VariableResolutionResamplerModel(nn.Module): + def __init__( + self, + in_dim, + out_dim, + spatial_conv_size, + temporal_conv_size, + config, + prefix: str = "", + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.config = config + self.spatial_conv_size = spatial_conv_size + self.temporal_conv_size = temporal_conv_size + self.use_temporal_conv = config.use_temporal_conv + + # compress 2d conv(picture) to 1d + self.spatial_dim = self.in_dim * self.spatial_conv_size * self.spatial_conv_size + # compress 3d conv(video) to 1d + self.temporal_dim = ( + self.in_dim + * self.spatial_conv_size + * self.spatial_conv_size + * self.temporal_conv_size + ) + + self.spatial_linear1 = ColumnParallelLinear( + self.spatial_dim, + self.spatial_dim, + bias=True, + gather_output=True, + quant_config=getattr(config, "quant_config", None), + prefix=f"{prefix}.spatial_linear1", + ) + + self.spatial_gelu = nn.GELU() + + self.spatial_linear2 = ColumnParallelLinear( + self.spatial_dim, + self.spatial_dim, + bias=True, + gather_output=True, + quant_config=getattr(config, "quant_config", None), + prefix=f"{prefix}.spatial_linear2", + ) + + self.spatial_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6) + + if self.use_temporal_conv: + self.temporal_linear1 = ColumnParallelLinear( + self.temporal_dim, + self.spatial_dim, + bias=True, + gather_output=True, + quant_config=getattr(config, "quant_config", None), + prefix=f"{prefix}.temporal_linear1", + ) + + self.temporal_gelu = nn.GELU() + + self.temporal_linear2 = ColumnParallelLinear( + self.spatial_dim, + self.spatial_dim, + bias=True, + gather_output=True, + quant_config=getattr(config, "quant_config", None), + prefix=f"{prefix}.temporal_linear2", + ) + + self.temporal_norm = nn.LayerNorm(self.spatial_dim, eps=1e-6) + + self.mlp = ColumnParallelLinear( + self.spatial_dim, + self.out_dim, + bias=True, + gather_output=True, + quant_config=getattr(config, "quant_config", None), + prefix=f"{prefix}.mlp", + ) + + self.after_norm = RMSNorm( + hidden_size=out_dim, eps=getattr(config, "rms_norm_eps", 1e-6) + ) + + def spatial_conv_reshape(self, x, spatial_conv_size): + S, C = x.shape + x = x.reshape([-1, C * (spatial_conv_size**2)]) + return x + + def forward(self, x, grid_thw): + def fwd_spatial(x): + x = self.spatial_conv_reshape(x, self.spatial_conv_size) + + x, _ = self.spatial_linear1(x) + x = self.spatial_gelu(x) + x, _ = self.spatial_linear2(x) + x = self.spatial_norm(x) + + return x + + def fwd_placeholder(x, grid_thw, to_tensor=False): + grid_thw_cpu = grid_thw.cpu().numpy() + grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:] + grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**2) + + tokens_per_img_or_vid = grid_thw_cpu.prod(-1) // (self.spatial_conv_size**2) + batch_offset = np.empty( + tokens_per_img_or_vid.size, dtype=tokens_per_img_or_vid.dtype + ) + batch_offset[0] = 0 + batch_offset[1:] = tokens_per_img_or_vid.cumsum()[:-1] + + slice_offsets = [] + for temporoal_size, spatial_size, b_offset in zip( + grid_t, grid_hw_after_conv, batch_offset + ): + for temp_offset in range(0, temporoal_size, 2): + slice_offsets.append( + np.arange( + b_offset + (temp_offset) * spatial_size, + b_offset + (temp_offset + 1) * spatial_size, + ) + ) + slice_offsets = torch.tensor(np.concatenate(slice_offsets, axis=-1)).to( + x.device + ) + + slice_offsets2 = [] + for temporoal_size, spatial_size, b_offset in zip( + grid_t, grid_hw_after_conv, batch_offset + ): + for temp_offset in range( + 1 if temporoal_size > 1 else 0, temporoal_size, 2 + ): + slice_offsets2.append( + np.arange( + b_offset + (temp_offset) * spatial_size, + b_offset + (temp_offset + 1) * spatial_size, + ) + ) + slice_offsets2 = torch.tensor(np.concatenate(slice_offsets2, axis=-1)).to( + x.device + ) + + x_timestep_1 = torch.index_select(x, dim=0, index=slice_offsets) + x_timestep_2 = torch.index_select(x, dim=0, index=slice_offsets2) + x = torch.concat([x_timestep_1, x_timestep_2], dim=-1) + return x + + def fwd_temporal(x): + x, _ = self.temporal_linear1(x) + x = self.temporal_gelu(x) + x, _ = self.temporal_linear2(x) + x = self.temporal_norm(x) + return x + + def fwd_mlp(x): + x, _ = self.mlp(x) + x = self.after_norm(x) + return x + + x = fwd_spatial(x) + if self.use_temporal_conv: + x = fwd_placeholder(x, grid_thw) + x = fwd_temporal(x) + x = fwd_mlp(x) + return x + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class Ernie4_5_VisionRotaryEmbedding(nn.Module): + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + self.inv_freq = 1.0 / theta ** ( + torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim + ) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + freqs = torch.outer(input=seq, vec2=self.inv_freq) + return freqs + + +class Ernie4_5_VisionTransformer(nn.Module): + + def __init__( + self, + vision_config: PretrainedConfig, + norm_eps: float = 1e-6, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + patch_size: int = vision_config.patch_size + spatial_merge_size: int = vision_config.spatial_merge_size + in_chans: int = vision_config.in_chans + hidden_size: int = vision_config.hidden_size + embed_dim: int = vision_config.embed_dim + depth: int = vision_config.depth + num_heads: int = vision_config.num_heads + mlp_ratio: float = vision_config.mlp_ratio + + self.spatial_merge_size = spatial_merge_size + + self.patch_embed = Ernie4_5_VisionPatchEmbed( + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ) + + norm_layer = partial(nn.LayerNorm, eps=norm_eps) + head_dim = embed_dim // num_heads + self.rotary_pos_emb = Ernie4_5_VisionRotaryEmbedding(head_dim // 2) + self.blocks = nn.ModuleList( + [ + Ernie4_5_VisionBlock( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{i}", prefix), + ) + for i in range(depth) + ] + ) + + self.ln = nn.LayerNorm(hidden_size, eps=1e-6) + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.blocks[0].mlp.fc2.weight.device + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + pos_ids = [] + for i in range(grid_thw.size(0)): + t, h, w = grid_thw[i].tolist() + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + hpos_ids = ( + hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + wpos_ids = ( + wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def forward( + self, + x: torch.Tensor, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + # patchify + x = x.to(device=self.device, dtype=self.dtype) + x = self.patch_embed(x) + + # compute position embedding + rotary_pos_emb = self.rot_pos_emb(grid_thw) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + # compute cu_seqlens + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum(dim=0, dtype=torch.int32) + cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens]) + + # transformers + x = x.unsqueeze(1) + for blk in self.blocks: + x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings) + + final_output = self.ln(x) + + if final_output.ndim == 3: + final_output = final_output.squeeze(dim=1) + + return final_output + + +cached_get_processor = lru_cache(get_processor) + + +class Ernie4_5_VLMoeForConditionalGeneration(nn.Module): + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.vision_model = Ernie4_5_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=add_prefix("vision_model", prefix), + ) + + self.model = Ernie4_5_VLMoeModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + + self.resampler_model = VariableResolutionResamplerModel( + self.config.pixel_hidden_size, + self.config.hidden_size, + self.config.spatial_conv_size, + self.config.temporal_conv_size, + config=self.config, + prefix=add_prefix("resampler_model", prefix), + ) + + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + + self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling + self.logits_processor = LogitsProcessor(config) + + if getattr(self.config, "im_patch_id", None): + visual_token_ids = [ + token_id + for token_id in [ + self.config.im_patch_id, + getattr(self.config, "image_start_token_id", None), + getattr(self.config, "image_end_token_id", None), + getattr(self.config, "video_start_token_id", None), + getattr(self.config, "video_end_token_id", None), + ] + if token_id is not None + ] + self._visual_token_ids_tensor_cache = torch.tensor( + visual_token_ids, dtype=torch.long + ) + else: + self._visual_token_ids_tensor_cache = None + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + pattern = MultiModalityDataPaddingPatternMultimodalTokens() + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def _vision_forward( + self, + pixel_values: torch.Tensor, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + if grid_thw is not None: + grid_thw = grid_thw[grid_thw > 0] + if grid_thw.numel() % 3 != 0: + raise ValueError( + f"grid_thw has {grid_thw.numel()} elements after filtering," + "which is not divisible by 3." + ) + grid_thw = grid_thw.reshape(-1, 3) + # example: [[1,64,64],[2,80,80]] -> [[1,64,64],[1,80,80],[1,80,80]] + grid_thw = F.pad( + torch.repeat_interleave(grid_thw[:, 1:], grid_thw[:, 0], 0), + [1, 0, 0, 0], + value=1, + ) + image_features = self.vision_model(pixel_values, grid_thw) + return image_features + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + # in qwen-vl, last dim is the same + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + self.vision_model.dtype + ) + image_grid_thw = torch.concat([item.image_grid_thw for item in items], dim=0) + assert pixel_values.dim() == 2, pixel_values.dim() + assert image_grid_thw.dim() == 2, image_grid_thw.dim() + image_feature = self._vision_forward(pixel_values, grid_thw=image_grid_thw) + image_embeds = self.resampler_model(image_feature, image_grid_thw) + return image_embeds + + def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + # in qwen-vl, last dim is the same + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + self.vision_model.dtype + ) + video_grid_thw = torch.concat([item.video_grid_thw for item in items], dim=0) + assert pixel_values.dim() == 2, pixel_values.dim() + assert video_grid_thw.dim() == 2, video_grid_thw.dim() + video_feature = self._vision_forward(pixel_values, grid_thw=video_grid_thw) + video_embeds = self.resampler_model(video_feature, video_grid_thw) + return video_embeds + + def _set_visual_token_mask( + self, input_ids: torch.Tensor, forward_batch: ForwardBatch + ) -> None: + """Set mask for visual tokens (image/video patches and delimiters).""" + if self._visual_token_ids_tensor_cache is None: + self.visual_token_mask = None + return + # Create tensor on the correct device + visual_token_ids_tensor = self._visual_token_ids_tensor_cache.to( + device=input_ids.device, + dtype=input_ids.dtype, + ) + + pad_values = [] + if hasattr(forward_batch, "mm_inputs") and forward_batch.mm_inputs is not None: + for mm_input in forward_batch.mm_inputs: + if mm_input is None: + continue + for item in mm_input.mm_items: + pad_values.append(item.pad_value) + placeholder_tensor = torch.as_tensor( + pad_values, + device=input_ids.device, + ) + pad_visual_token_ids_tensor = torch.cat( + [visual_token_ids_tensor, placeholder_tensor], dim=0 + ) + self.visual_token_mask = torch.isin( + input_ids, pad_visual_token_ids_tensor + ).reshape(-1, 1) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def should_apply_lora(self, module_name: str) -> bool: + # skip vision_model + return not module_name.startswith("vision_model") + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + ): + """Run forward pass for Ernie45-VL. + + Args: + input_ids: Flattened (concatenated) input_ids corresponding to a + batch. + positions: Flattened (concatenated) position ids corresponding to a + batch. + **NOTE**: If mrope is enabled (default setting for Qwen2-VL + opensource models), the shape will be `(3, seq_len)`, + otherwise it will be `(seq_len,). + (Use input_metadata.mrope_positions to replace it) + """ + if self.is_mrope_enabled: + positions = forward_batch.mrope_positions + + if not ( + forward_batch.forward_mode.is_decode() + or not forward_batch.contains_image_inputs() + ): + if self.is_mrope_enabled: + assert positions.ndim == 2 and positions.size(0) == 3, ( + "multimodal section rotary embedding requires " + f"(3, seq_len) positions, but got {positions.size()}" + ) + + self._set_visual_token_mask(input_ids, forward_batch) + + assert ( + input_ids.numel() == positions.shape[-1] + ), f"input_ids {input_ids.shape} and position_ids {positions.shape} should have the same length" + + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + multimodal_model=self, + positions=positions, + visual_token_mask=self.visual_token_mask, + ) + + self.visual_token_mask = None + + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "up_proj", 1), + ("gate_up_proj", "gate_proj", 0), + ] + + # resampler_weight_mappings + resampler_weight_mapping = { + "spatial_linear.0.": "spatial_linear1.", + "spatial_linear.2.": "spatial_linear2.", + "spatial_linear.3.": "spatial_norm.", + "temporal_linear.0.": "temporal_linear1.", + "temporal_linear.2.": "temporal_linear2.", + "temporal_linear.3.": "temporal_norm.", + } + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=max(self.config.moe_num_experts), + ) + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if "vision_model" in name: + # adapt to VisionAttention + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + if name.startswith("model.resampler_model"): + name = name.replace("model.resampler_model", "resampler_model") + + for ( + old_weight_name, + new_weight_name, + ) in resampler_weight_mapping.items(): + if old_weight_name in name: + name = name.replace(old_weight_name, new_weight_name, 1) + break + + # Distinguish between vision experts and text experts + if "mlp.experts" in name: + moe_offset = int(name.split(".")[-3]) + vision_expert_start_idx = self.config.moe_num_experts[0] + is_text_expert = moe_offset <= vision_expert_start_idx - 1 + if is_text_expert: + name = name.replace(".experts.", ".text_experts.") + else: + name = name.replace( + f".experts.{moe_offset}", + f".vision_experts.{moe_offset - vision_expert_start_idx}", + ) + + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + + # Distinguish between vision experts and text experts + moe_offset = int(name.split(".")[-3]) + is_text_expert = moe_offset <= self.config.moe_num_experts[0] - 1 + + name = name.replace(weight_name, param_name) + if is_text_expert: + name = name.replace(".experts.", ".text_experts.") + else: + name = name.replace(".experts.", ".vision_experts.") + + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + else: + logger.warning(f"Parameter {name} not found in params_dict") + break + else: + # Distinguish between vision expert gate + # and text expert gate + if name.endswith("mlp.gate.weight"): + name = name.replace("gate.weight", "text_experts_gate.weight") + loaded_weight = loaded_weight.T + elif name.endswith("mlp.gate.weight_1"): + name = name.replace( + "gate.weight_1", "vision_experts_gate.weight" + ) + loaded_weight = loaded_weight.T + + if "e_score_correction_bias" in name: + name = name.replace(".moe_statics.", ".") + + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") + + +EntryClass = [Ernie4_5_VLMoeForConditionalGeneration] diff --git a/python/sglang/srt/multimodal/processors/ernie45_vl.py b/python/sglang/srt/multimodal/processors/ernie45_vl.py new file mode 100644 index 000000000..7ec25015a --- /dev/null +++ b/python/sglang/srt/multimodal/processors/ernie45_vl.py @@ -0,0 +1,417 @@ +import math +import os +from typing import List, Union + +import numpy as np +import torch +import torchvision +from PIL import Image +from torchvision.transforms import InterpolationMode +from transformers import BaseImageProcessorFast + +from sglang.srt.environ import envs +from sglang.srt.layers.rotary_embedding import MRotaryEmbedding +from sglang.srt.models.ernie45_vl import Ernie4_5_VLMoeForConditionalGeneration +from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor as SGLangBaseProcessor, +) +from sglang.srt.multimodal.processors.base_processor import MultimodalSpecialTokens +from sglang.srt.utils import get_bool_env_var, is_npu, logger + +_is_npu = is_npu() + +SGL_USE_CUDA_IPC = get_bool_env_var("SGLANG_USE_CUDA_IPC_TRANSPORT") + + +IMAGE_FACTOR = 28 +MIN_PIXELS = 4 * 28 * 28 +# MAX_PIXELS = envs.SGLANG_IMAGE_MAX_PIXELS.get() +MAX_PIXELS = 16384 * 28 * 28 +MAX_RATIO = 200 +RESIZE_RESAMPLE = getattr(Image, envs.SGLANG_RESIZE_RESAMPLE.get(), None) +if envs.SGLANG_RESIZE_RESAMPLE.is_set() and RESIZE_RESAMPLE is None: + logger.warning( + f"Invalid RESIZE_RESAMPLE value: '{envs.SGLANG_RESIZE_RESAMPLE.get()}'. " + f"Ignoring and using default." + ) +VIDEO_TOTAL_PIXELS = int( + float(os.environ.get("VIDEO_MAX_PIXELS", 128000 * 28 * 28 * 0.9)) +) + +VIDEO_MIN_PIXELS = 299 * 28 * 28 +VIDEO_MAX_PIXELS = 1196 * 28 * 28 +FRAME_FACTOR = 2 +FPS = 2.0 +FPS_MIN_FRAMES = 16 +FPS_MAX_FRAMES = 180 + + +def smart_resize( + height: int, + width: int, + factor: int = IMAGE_FACTOR, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS, +): + if max(height, width) / min(height, width) > MAX_RATIO: + if height > width: + new_width = max(factor, round_by_factor(width, factor)) + new_height = floor_by_factor(new_width * MAX_RATIO, factor) + else: + new_height = max(factor, round_by_factor(height, factor)) + new_width = floor_by_factor(new_height * MAX_RATIO, factor) + + height = new_height + width = new_width + + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + + if min_pixels > h_bar * w_bar or h_bar * w_bar > max_pixels: + raise ValueError(f"encounter invalid h_bar: {h_bar}, w_bar: {w_bar}") + + return h_bar, w_bar + + +def resize_image( + image, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS, + size_factor: int = IMAGE_FACTOR, +) -> Image.Image: + width, height = image.size + min_pixels = min_pixels + max_pixels = max_pixels + resized_height, resized_width = smart_resize( + height, + width, + factor=size_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height), resample=RESIZE_RESAMPLE) + return image + + +def round_by_factor(number: int | float, factor: int) -> int: + return round(number / factor) * factor + + +def ceil_by_factor(number: int | float, factor: int) -> int: + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int | float, factor: int) -> int: + return math.floor(number / factor) * factor + + +async def resize_image_async( + image, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS, + size_factor: int = IMAGE_FACTOR, +): + return resize_image(image, min_pixels, max_pixels, size_factor) + + +def smart_nframes( + ele: dict, + total_frames: int, + video_fps: int | float, +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ( + "fps" in ele and "nframes" in ele + ), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) + max_frames = floor_by_factor( + ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR + ) + nframes = total_frames / video_fps * fps + if nframes > total_frames: + logger.warning( + f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]" + ) + nframes = min(min(max(nframes, min_frames), max_frames), total_frames) + nframes = floor_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError( + f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}." + ) + return nframes + + +# process video, qwen-specific +async def preprocess_video( + vr, + image_factor: int = IMAGE_FACTOR, +) -> torch.Tensor: + + total_frames, video_fps = len(vr), vr.get_avg_fps() + nframes = smart_nframes({}, total_frames=total_frames, video_fps=video_fps) + idx = np.linspace(0, total_frames - 1, num=nframes, dtype=np.int64) + idx = np.unique(idx) + video_np = vr.get_batch(idx).asnumpy() + video = torch.from_numpy(video_np).pin_memory() + video = video.permute(0, 3, 1, 2) # Convert to TCHW format + nframes, _, height, width = video.shape + min_pixels = VIDEO_MIN_PIXELS + total_pixels = VIDEO_TOTAL_PIXELS + max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05), + ) + + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + video = torchvision.transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BILINEAR, + ) + + video = video.permute(0, 2, 3, 1) + video = video.pin_memory() + video_metadata = { + "fps": video_fps, + "duration": total_frames / video_fps, + "total_num_frames": total_frames, + "frames_indices": idx, + "video_backend": "torchvision", + } + + return video, video_metadata + + +# Compatible with Ernie-VL Series +class Ernie4_5_VLImageProcessor(SGLangBaseProcessor): + models = [Ernie4_5_VLMoeForConditionalGeneration] + + def __init__(self, hf_config, server_args, _processor, *args, **kwargs): + super().__init__(hf_config, server_args, _processor, *args, **kwargs) + self.hf_config = hf_config + self.model_type = hf_config.model_type + self.image_start_token_id = hf_config.image_start_token_id + self.image_end_token_id = hf_config.image_end_token_id + self.video_start_token_id = hf_config.video_start_token_id + self.video_end_token_id = hf_config.video_end_token_id + + self.IMAGE_FACTOR = 28 + self.MIN_PIXELS = 4 * 28 * 28 + self.MAX_PIXELS = 16384 * 28 * 28 + self.MAX_RATIO = 200 + self.mm_tokens = MultimodalSpecialTokens( + image_token="<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>", + video_token="<|VIDEO_START|><|video@placeholder|><|VIDEO_END|>", + image_token_id=hf_config.im_patch_id, + video_token_id=hf_config.im_patch_id, # image and video use the same token_id + ).build(_processor) + + self.tokenizer = self._processor.tokenizer + self.image_processor = self._processor.image_processor + + def _pixel_values_norm( + self, + pixel_values: torch.Tensor, + mm_kwargs: object, + ) -> torch.Tensor: + hf_config = self.hf_config + vision_config = hf_config.vision_config + image_processor = self.image_processor + image_mean_tensor = torch.tensor( + image_processor.image_mean, dtype=torch.float32 + ).reshape([1, 3, 1, 1]) + image_std_tensor = torch.tensor( + image_processor.image_std, dtype=torch.float32 + ).reshape([1, 3, 1, 1]) + rescale_factor = torch.tensor( + image_processor.rescale_factor, dtype=torch.float32 + ) + patch_size_squared = vision_config.patch_size**2 + + image_mean_tensor = image_mean_tensor.squeeze([-2, -1]).repeat_interleave( + patch_size_squared, -1 + ) + image_std_tensor = image_std_tensor.squeeze([-2, -1]).repeat_interleave( + patch_size_squared, -1 + ) + + if not image_mean_tensor.is_contiguous(): + image_mean_tensor = image_mean_tensor.contiguous() + if not image_std_tensor.is_contiguous(): + image_std_tensor = image_std_tensor.contiguous() + + pixel_values = ( + rescale_factor * pixel_values.to(torch.float32) - image_mean_tensor + ) / image_std_tensor + pixel_values = pixel_values.to(hf_config.dtype) + return pixel_values + + def process_mm_data( + self, input_text, images=None, videos=None, audios=None, **kwargs + ) -> dict: + """ + process multimodal data with transformers AutoProcessor + """ + if images: + kwargs["images"] = images + if videos: + kwargs["videos"] = videos + + processor = self._processor + if ( + hasattr(processor, "image_processor") + and isinstance(processor.image_processor, BaseImageProcessorFast) + and not self.server_args.disable_fast_image_processor + ): + if not _is_npu: + kwargs["device"] = "cuda" + + result = processor.__call__( + text=[input_text], + padding=True, + return_tensors="pt", + **kwargs, + ) + + # Divide the processor_output into two modalities: image and video. + if result is not None: + pixel_values = result["images"] + if pixel_values is not None: + result["images"] = self._pixel_values_norm(pixel_values, kwargs) + for key in list(result.keys()): + if result[key] is None: + del result[key] + continue + if key == "grid_thw": + grid_thw = result["grid_thw"] + pixel_values_all = result["images"] + # Identify elements where the first + # dimension is greater than 1 and + # treat them as the video modality + mask = grid_thw[:, 0] > 1 + result["video_grid_thw"] = grid_thw[mask] + result["image_grid_thw"] = grid_thw[~mask] + image_patch_num = result["image_grid_thw"].prod(dim=1).sum() + result["pixel_values"] = pixel_values_all[:image_patch_num] + result["pixel_values_videos"] = pixel_values_all[image_patch_num:] + del result["images"] + del result["grid_thw"] + + # del empty result + if result["image_grid_thw"].numel() == 0: + del result["image_grid_thw"] + if result["pixel_values"].numel() == 0: + del result["pixel_values"] + if result["video_grid_thw"].numel() == 0: + del result["video_grid_thw"] + if result["pixel_values_videos"].numel() == 0: + del result["pixel_values_videos"] + + if not self.server_args.keep_mm_feature_on_device: + # move feature tensors to cpu + for feature_name in self.FEATURE_NAMES: + if SGL_USE_CUDA_IPC: + pass + else: + if feature_name in result and isinstance( + result[feature_name], torch.Tensor + ): + result[feature_name] = result[feature_name].to("cpu") + + return result + + async def process_mm_data_async( + self, + image_data: List[Union[str, bytes]], + input_text, + request_obj, + *args, + **kwargs, + ): + base_output = self.load_mm_data( + prompt=input_text, + image_data=image_data, + video_data=request_obj.video_data, + audio_data=request_obj.audio_data, + multimodal_tokens=self.mm_tokens, + ) + + # resize images if they are raw Image objects + resized_images = [] + if base_output.images and isinstance(base_output.images[0], Image.Image): + for image in base_output.images: + resized_image = resize_image(image) + resized_images.append(resized_image) + base_output.images = resized_images + + if base_output.videos: + videos_processed = [ + await preprocess_video(video) for video in base_output.videos + ] + base_output.videos, _ = map(list, zip(*videos_processed)) + + mm_items, input_ids, ret = self.process_and_combine_mm_data( + base_output, self.mm_tokens + ) + + input_ids = input_ids.flatten() + + mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index_ernie45( + input_ids=input_ids.unsqueeze(0), + hf_config=self.hf_config, + image_grid_thw=getattr(ret, "image_grid_thw", None), + video_grid_thw=getattr(ret, "video_grid_thw", None), + ) + mrope_positions = mrope_positions.squeeze(1) + + assert ( + input_ids.shape[0] == mrope_positions.shape[-1] + ), "input_ids and mrope_positions should have the same length" + + mm_inputs = { + "input_ids": input_ids.tolist(), + "mm_items": mm_items, + "im_start_id": self.image_start_token_id, + "im_end_id": self.image_end_token_id, + "im_token_id": self.mm_tokens.image_token_id, + "video_token_id": self.mm_tokens.video_token_id, + "mrope_positions": mrope_positions, + "mrope_position_delta": mrope_position_delta, + } + + return mm_inputs