From 3e7ecb78a60f8e1d889cfe25c88006577783d903 Mon Sep 17 00:00:00 2001 From: RunningLeon Date: Thu, 5 Feb 2026 00:22:44 +0800 Subject: [PATCH] model: support interns1-pro (#18145) Co-authored-by: Ke Bao --- python/sglang/srt/configs/model_config.py | 1 + python/sglang/srt/layers/rotary_embedding.py | 199 ++++++++++++++ python/sglang/srt/models/interns1pro.py | 252 ++++++++++++++++++ python/sglang/srt/models/qwen3_vl_moe.py | 10 +- .../srt/multimodal/processors/interns1pro.py | 118 ++++++++ 5 files changed, 578 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/models/interns1pro.py create mode 100644 python/sglang/srt/multimodal/processors/interns1pro.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 64b9f6fd3..e76c15705 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1197,6 +1197,7 @@ multimodal_model_archs = [ "KimiVLForConditionalGeneration", "InternVLChatModel", "InternS1ForConditionalGeneration", + "InternS1ProForConditionalGeneration", "Phi4MMForCausalLM", "Step3VLForConditionalGeneration", "POINTSV15ChatModel", diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index da3d3ecb0..311df6d01 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn +import torch.nn.functional as F import triton import triton.language as tl @@ -801,6 +802,191 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: return 0.1 * mscale * math.log(scale) + 1.0 +class FourierRotaryEmbedding(nn.Module): + """Fourier RotaryEmbedding extended.""" + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + dtype: torch.dtype, + num_kv_heads: int, + *, + fope_init_factor: float = 0.1, + fope_sep_head: bool = True, + num_inv_freq: int = None, + device: Optional[str] = "cuda", + ) -> None: + self.fope_init_factor = fope_init_factor + self.fope_sep_head = fope_sep_head + self.num_inv_freq = num_inv_freq + self.num_kv_heads = num_kv_heads + self.device = device + + super().__init__() + self.head_size = head_size + self.rotary_dim = rotary_dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.is_neox_style = is_neox_style + self.dtype = dtype + + self.fope_init_factor = fope_init_factor + self.fope_sep_head = fope_sep_head + self.num_inv_freq = num_inv_freq + self.num_kv_heads = num_kv_heads + + self.inv_freq: torch.Tensor + self.register_buffer( + "inv_freq", self._compute_inv_freq(self.base), persistent=False + ) + self.input_dim = self.inv_freq.shape[-1] + self.output_dim = self.inv_freq.shape[-1] + self.cos_coef = nn.Parameter( + torch.empty( + self.num_kv_heads, self.input_dim, self.output_dim, dtype=torch.float32 + ), + requires_grad=False, + ) + self.sin_coef = nn.Parameter( + torch.empty( + self.num_kv_heads, self.input_dim, self.output_dim, dtype=torch.float32 + ), + requires_grad=False, + ) + self.cos_sin_cache: torch.Tensor + self.register_buffer( + "cos_sin_cache", self._compute_cos_sin_cache(), persistent=False + ) + # update cos_sin_cache after update weights + self.update_buffer = False + + def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor: + """Compute the inverse frequency.""" + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.int64).to( + device=self.device, dtype=torch.float + ) + / self.rotary_dim + ) + ) + + assert ( + inv_freq[:-1] > inv_freq[1:] + ).all(), "Expected inv_freq to be in decreasing order" + + inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) + if self.num_inv_freq is not None: + inv_freq_idx_selected[self.num_inv_freq :] = False + else: + inv_freq_idx_selected = inv_freq > ( + 2.0 * torch.pi / self.max_position_embeddings + ) + + inv_freq = inv_freq[inv_freq_idx_selected] + return inv_freq + + def _compute_cos_sin_cache(self) -> torch.Tensor: + """Compute the cos and sin cache.""" + + t = torch.arange( + self.max_position_embeddings, dtype=torch.float, device=self.device + ) + + freqs = torch.einsum("i,j -> ij", t, self.inv_freq) + if self.fope_sep_head: + pos_cos = freqs.cos().unsqueeze(0).expand(self.num_kv_heads, -1, -1) + pos_sin = freqs.sin().unsqueeze(0).expand(self.num_kv_heads, -1, -1) + else: + pos_cos = freqs.cos() + pos_sin = freqs.sin() + + if self.fope_sep_head: + sin = torch.einsum("htD, hDd -> thd", pos_sin, self.sin_coef.float()) + cos = torch.einsum("htD, hDd -> thd", pos_cos, self.cos_coef.float()) + else: + sin = torch.einsum("tD, Dd -> td", pos_sin, self.sin_coef.float()) + cos = torch.einsum("tD, Dd -> td", pos_cos, self.cos_coef.float()) + + sin = F.pad( + input=sin, + pad=(0, self.head_size // 2 - sin.size(-1)), + mode="constant", + value=1, + ) + cos = F.pad( + input=cos, + pad=(0, self.head_size // 2 - cos.size(-1)), + mode="constant", + value=1, + ) + + sin = torch.cat((sin, sin), dim=-1) + cos = torch.cat((cos, cos), dim=-1) + + cache = torch.cat((cos, sin), dim=-1) + return cache + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if not self.update_buffer: + self.cos_sin_cache = self._compute_cos_sin_cache() + self.update_buffer = True + + query = query.unflatten(-1, (-1, self.head_size)) + key = key.unflatten(-1, (-1, self.head_size)) + positions_with_offsets = ( + torch.add(positions, offsets) if offsets is not None else positions + ) + cos_sin = torch.index_select(self.cos_sin_cache, 0, positions_with_offsets).to( + dtype=query.dtype + ) + cos, sin = cos_sin.chunk(2, dim=-1) + + assert ( + query.dim() == key.dim() == 3 + ), "Expected query key (seq_len, heads, head_dim)" + assert cos.dim() <= 3 and sin.dim() <= 3 + + need_reshape = False + if cos.dim() == 3: + # for fope + need_reshape = True + query_shape = query.shape + key_shape = key.shape + cos = cos.flatten(0, 1) + sin = sin.flatten(0, 1) + seq_len = cos.size(0) + query = query.reshape(seq_len, -1, query.size(-1)) + key = key.reshape(seq_len, -1, key.size(-1)) + + query, key = apply_rotary_pos_emb_native(query, key, cos, sin) + + if need_reshape: + query = query.reshape(query_shape) + key = key.reshape(key_shape) + return query.flatten(-2), key.flatten(-2) + + def extra_repr(self) -> str: + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" + s += f", max_position_embeddings={self.max_position_embeddings}" + s += f", base={self.base}, is_neox_style={self.is_neox_style}" + s += f", fope_init_factor={self.fope_init_factor}, fope_sep_head={self.fope_sep_head}" + s += f", num_inv_freq={self.num_inv_freq}, num_kv_heads={self.num_kv_heads}" + return s + + class DeepseekScalingRotaryEmbedding(RotaryEmbedding): """RotaryEmbedding extended with YaRN method. @@ -2901,6 +3087,19 @@ def get_rope( mrope_section=rope_scaling["mrope_section"], mrope_interleaved=rope_scaling.get("mrope_interleaved", False), ) + elif rope_scaling.get("use_fope", False): + rotary_emb = FourierRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + num_kv_heads=rope_scaling["num_kv_heads"], + fope_init_factor=rope_scaling.get("fope_init_factor", 0.1), + fope_sep_head=rope_scaling.get("fope_sep_head", True), + num_inv_freq=rope_scaling.get("num_inv_freq", None), + ) else: rotary_emb = RotaryEmbedding( head_size, diff --git a/python/sglang/srt/models/interns1pro.py b/python/sglang/srt/models/interns1pro.py new file mode 100644 index 000000000..b22ff20ec --- /dev/null +++ b/python/sglang/srt/models/interns1pro.py @@ -0,0 +1,252 @@ +import functools +import logging +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from transformers import PretrainedConfig + +from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.rotary_embedding import get_rope +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.qwen3_moe import Qwen3MoeAttention, Qwen3MoeDecoderLayer +from sglang.srt.models.qwen3_vl_moe import ( + Qwen3MoeLLMModel, + Qwen3VLMoeForConditionalGeneration, +) +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class InternS1ProTextAttention(Qwen3MoeAttention): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 1000000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 32768, + **kwargs, + ) -> None: + super().__init__( + hidden_size, + num_heads, + num_kv_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + **kwargs, + ) + # for fope + fope_keys = {"fope_init_factor", "fope_sep_head", "num_inv_freq"} + use_fope = any(rope_scaling.get(key) is not None for key in fope_keys) + if use_fope: + rope_scaling["use_fope"] = True + rope_scaling["num_kv_heads"] = self.num_kv_heads + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.compatible_with_fused_kv_buffer = False + self.use_fused_qk_norm_rope = False + self._used_fused_qk_norm_rope_last_call = False + + def forward_prepare_npu( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ): + raise NotImplementedError() + + +class InternS1ProTextDecoderLayer(Qwen3MoeDecoderLayer): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__( + config, + layer_id, + quant_config=quant_config, + prefix=prefix, + alt_stream=alt_stream, + ) + + rope_theta = getattr(config, "rope_theta", 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + rms_norm_eps = config.rms_norm_eps + attention_bias = config.attention_bias + + self.self_attn = InternS1ProTextAttention( + hidden_size=self.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, + max_position_embeddings=max_position_embeddings, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + attention_bias=attention_bias, + config=config, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, + ) + # update with group router + self.router_n_groups = getattr(config, "router_n_groups", -1) + if self.router_n_groups > 0: + assert ( + config.num_experts_per_tok % self.router_n_groups == 0 + ), f"{config.num_experts_per_tok} cannot be divided by {self.router_n_groups}" + self.mlp.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + use_grouped_topk=False, + layer_id=layer_id, + custom_routing_function=self._custom_routing_function, + ) + + @staticmethod + @functools.lru_cache + def get_group_offsets(router_n_groups: int, group_size: int, device: str): + group_offsets = ( + torch.arange(router_n_groups, device=device) * group_size + ).view( + 1, -1, 1 + ) # [1, n_groups, 1] + return group_offsets + + def _custom_routing_function( + self, + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> torch.Tensor: + """Group router""" + routing_weights = torch.softmax(gating_output, dim=-1, dtype=torch.float32) + if self.router_n_groups > 0: + assert ( + routing_weights.shape[-1] % self.router_n_groups == 0 + ), f"{routing_weights.shape[-1]} cannot be divided by {self.router_n_groups}" + per_group_top_k = topk // self.router_n_groups + group_size = routing_weights.shape[-1] // self.router_n_groups + group_offsets = self.get_group_offsets( + self.router_n_groups, group_size, routing_weights.device + ) + routing_weights = routing_weights.unflatten( + -1, (self.router_n_groups, group_size) + ) + topk_weights, topk_ids = torch.topk( + routing_weights, per_group_top_k, dim=-1 + ) + topk_ids = (topk_ids + group_offsets).flatten(-2, -1) + topk_weights = topk_weights.flatten(-2, -1) + else: + topk_weights, topk_ids = torch.topk(routing_weights, topk, dim=-1) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + return topk_weights, topk_ids + + +class InternS1ProTextModel(Qwen3MoeLLMModel): + def __init__( + self, + *, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + decoder_layer_type=InternS1ProTextDecoderLayer, + prefix: str = "", + ): + super().__init__( + config=config, + quant_config=quant_config, + prefix=prefix, + decoder_layer_type=decoder_layer_type, + ) + + +class InternS1ProForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + language_model_cls=InternS1ProTextModel, + ) -> None: + # deal with no deepstack + if not hasattr(config.vision_config, "deepstack_visual_indexes"): + config.vision_config.deepstack_visual_indexes = [] + + super().__init__( + config, + quant_config=quant_config, + prefix=prefix, + language_model_cls=language_model_cls, + ) + + # disable deepstack + if len(config.vision_config.deepstack_visual_indexes) == 0: + self.use_deepstack = {} + + def _load_fope_weights(self, name: str, loaded_weight: torch.Tensor, params_dict): + """load fope weights""" + attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_attention_tp_rank() + + num_key_value_heads = loaded_weight.size(0) + # replicate head if necessary + if num_key_value_heads < attn_tp_size: + n_replicate = attn_tp_size // num_key_value_heads + attn_tp_size = num_key_value_heads + attn_tp_rank = attn_tp_rank // n_replicate + loaded_weight = loaded_weight.chunk(attn_tp_size, dim=0)[attn_tp_rank] + + # rotary_emb is shared cross layers + param_name = name.replace(".rotary_emb.", ".layers.0.self_attn.rotary_emb.") + assert param_name in params_dict + param = params_dict[param_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + """load weights""" + # Cache params_dict to avoid repeated expensive traversal of model parameters + if not hasattr(self, "_cached_params_dict"): + self._cached_params_dict = dict(self.named_parameters()) + params_dict = self._cached_params_dict + other_weights = dict() + for name, loaded_weight in weights: + if "sin_coef" in name or "cos_coef" in name: + name = name.replace(r"model.language_model.", r"model.") + self._load_fope_weights(name, loaded_weight, params_dict) + else: + other_weights[name] = loaded_weight + + super().load_weights(other_weights.items()) + + +EntryClass = InternS1ProForConditionalGeneration diff --git a/python/sglang/srt/models/qwen3_vl_moe.py b/python/sglang/srt/models/qwen3_vl_moe.py index ed64010bf..671566c46 100644 --- a/python/sglang/srt/models/qwen3_vl_moe.py +++ b/python/sglang/srt/models/qwen3_vl_moe.py @@ -27,7 +27,7 @@ from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.models.qwen3_moe import Qwen3MoeModel +from sglang.srt.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration from sglang.srt.utils.hf_transformers_utils import get_processor @@ -43,8 +43,14 @@ class Qwen3MoeLLMModel(Qwen3MoeModel): config: Qwen3VLMoeTextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", + decoder_layer_type=Qwen3MoeDecoderLayer, ): - super().__init__(config=config, quant_config=quant_config, prefix=prefix) + super().__init__( + config=config, + quant_config=quant_config, + prefix=prefix, + decoder_layer_type=decoder_layer_type, + ) self.hidden_size = config.hidden_size # Currently, we use 3 as len(config.vision_config.deepstack_visual_indexes) is not directly accessible here. # This approach follows the original implementation. diff --git a/python/sglang/srt/multimodal/processors/interns1pro.py b/python/sglang/srt/multimodal/processors/interns1pro.py new file mode 100644 index 000000000..d448dc1cb --- /dev/null +++ b/python/sglang/srt/multimodal/processors/interns1pro.py @@ -0,0 +1,118 @@ +import time +from typing import List, Union + +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +from sglang.srt.models.interns1pro import InternS1ProForConditionalGeneration +from sglang.srt.multimodal.processors.qwen_vl import ( + QwenVLImageProcessor, + preprocess_video, +) +from sglang.utils import logger + + +class InternS1_1ImageProcessor(QwenVLImageProcessor): + models = [ + InternS1ProForConditionalGeneration, + ] + + def get_mm_data(self, prompt, embeddings, img_grid_thw): + input_ids, offsets = self.build_input_ids(prompt, img_grid_thw) + + mm_items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + offsets=offsets, + precomputed_embeddings=embeddings, + ) + ] + + return { + "input_ids": input_ids, + "mm_items": mm_items, + "im_start_id": self.IM_START_TOKEN_ID, + "im_end_id": self.IM_END_TOKEN_ID, + "im_token_id": self.mm_tokens.image_token_id, + "video_token_id": self.mm_tokens.video_token_id, + "audio_token_id": self.mm_tokens.audio_token_id, + } + + async def process_mm_data_async( + self, + image_data: List[Union[str, bytes]], + input_text, + request_obj, + *args, + **kwargs, + ): + entry_time = time.perf_counter() + 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, + ) + load_time = time.perf_counter() + rid = getattr(request_obj, "rid", "anonymous_rid") + + video_metadata = None + if base_output.videos: + videos_processed = [ + await preprocess_video(video, video_config=self.video_config) + for video in base_output.videos + ] + base_output.videos, video_metadata = map(list, zip(*videos_processed)) + + preprocess_time = time.perf_counter() + + mm_items, input_ids, ret = self.process_and_combine_mm_data( + base_output, + self.mm_tokens, + video_metadata=video_metadata, + do_sample_frames=False, + ) + + second_per_grid_ts = getattr(ret, "second_per_grid_ts", None) + if second_per_grid_ts is None: + second_per_grid_ts = getattr(ret, "video_second_per_grid", None) + + process_time = time.perf_counter() + + input_ids = input_ids.flatten() + + image_grid_thw = None + if hasattr(ret, "image_grid_thw"): + image_grid_thw = ret.image_grid_thw + + if image_grid_thw is None and image_data and isinstance(image_data[0], dict): + image_grid_thw = image_data[0].get("image_grid_thw") + + video_grid_thw = None + if hasattr(ret, "video_grid_thw"): + video_grid_thw = ret.video_grid_thw + + if video_grid_thw is None and request_obj.video_data: + first_video = request_obj.video_data[0] + if isinstance(first_video, dict): + video_grid_thw = first_video.get("video_grid_thw") + + get_rope_index_time = time.perf_counter() + + logger.debug( + f"[QwenVLProcessor Perf] {rid=}, " + f"load_time: {(load_time - entry_time) * 1000:.2f} ms, " + f"preprocess_time: {(preprocess_time - load_time) * 1000:.2f} ms, " + f"process_time: {(process_time - preprocess_time) * 1000:.2f} ms, " + f"get_rope_index_time: {(get_rope_index_time - process_time) * 1000:.2f} ms, " + f"total_time: {(get_rope_index_time - entry_time) * 1000:.2f} ms" + ) + + return { + "input_ids": input_ids.tolist(), + "mm_items": mm_items, + "im_start_id": self.vision_start_token_id, + "im_end_id": self.vision_end_token_id, + "im_token_id": self.mm_tokens.image_token_id, + "video_token_id": self.mm_tokens.video_token_id, + "audio_token_id": self.mm_tokens.audio_token_id, + }