model: support interns1-pro (#18145)

Co-authored-by: Ke Bao <ispobaoke@gmail.com>

python/sglang/srt/configs/model_config.py

@@ -1197,6 +1197,7 @@ multimodal_model_archs = [
     "KimiVLForConditionalGeneration",
     "InternVLChatModel",
     "InternS1ForConditionalGeneration",
+    "InternS1ProForConditionalGeneration",
     "Phi4MMForCausalLM",
     "Step3VLForConditionalGeneration",
     "POINTSV15ChatModel",

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,

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

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.

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,
+        }