@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a3acfa1071
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 00:00 UTC

[Models] Intern-S1-Pro (#33636)

Signed-off-by: zxy <zhou0493@e.ntu.edu.sg> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

zxy committed Feb 3, 2026 at 21:49 UTC a3acfa10719a931111caccd08ef19f1551b2fe1e
11 files changed +942 -11
docs/models/supported_models.md
+1
@@ -689,6 +689,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
689 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | |
690 | `IsaacForConditionalGeneration` | Isaac | T + I<sup>+</sup> | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ |
691 | `InternS1ForConditionalGeneration` | Intern-S1 | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |
692 +| `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ |
693 | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
694 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
695 | `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
examples/offline_inference/vision_language.py
+35
@@ -842,6 +842,40 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData:
842 )
843
844
845 +# Intern-S1-Pro
846 +def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData:
847 + model_name = "internlm/Intern-S1-Pro"
848 +
849 + engine_args = EngineArgs(
850 + model=model_name,
851 + trust_remote_code=True,
852 + max_model_len=8192,
853 + max_num_seqs=2,
854 + limit_mm_per_prompt={modality: 1},
855 + enforce_eager=True,
856 + tensor_parallel_size=4,
857 + )
858 +
859 + if modality == "image":
860 + placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
861 + elif modality == "video":
862 + placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
863 +
864 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
865 + messages = [
866 + [{"role": "user", "content": f"{placeholder}\n{question}"}]
867 + for question in questions
868 + ]
869 + prompts = tokenizer.apply_chat_template(
870 + messages, tokenize=False, add_generation_prompt=True
871 + )
872 +
873 + return ModelRequestData(
874 + engine_args=engine_args,
875 + prompts=prompts,
876 + )
877 +
878 +
879 # InternVL
880 def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
881 model_name = "OpenGVLab/InternVL3-2B"
@@ -2130,6 +2164,7 @@ model_example_map = {
2164 "hyperclovax_seed_vision": run_hyperclovax_seed_vision,
2165 "idefics3": run_idefics3,
2166 "interns1": run_interns1,
2167 + "interns1_pro": run_interns1_pro,
2168 "internvl_chat": run_internvl,
2169 "kanana_v": run_kanana_v,
2170 "keye_vl": run_keye_vl,
tests/models/registry.py
+6
@@ -755,6 +755,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
755 "InternS1ForConditionalGeneration": _HfExamplesInfo(
756 "internlm/Intern-S1", trust_remote_code=True
757 ),
758 + "InternS1ProForConditionalGeneration": _HfExamplesInfo(
759 + "internlm/Intern-S1-Pro",
760 + trust_remote_code=True,
761 + min_transformers_version="5.0.0",
762 + is_available_online=False,
763 + ),
764 "InternVLChatModel": _HfExamplesInfo(
765 "OpenGVLab/InternVL2-1B",
766 extras={
vllm/model_executor/layers/rotary_embedding/__init__.py
+23
@@ -11,6 +11,7 @@ from .deepseek_scaling_rope import DeepseekScalingRotaryEmbedding
11 from .dual_chunk_rope import DualChunkRotaryEmbedding
12 from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding
13 from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding
14 +from .fope import FourierRotaryEmbedding
15 from .linear_scaling_rope import LinearScalingRotaryEmbedding
16 from .llama3_rope import Llama3RotaryEmbedding
17 from .llama4_vision_rope import Llama4VisionRotaryEmbedding
@@ -102,6 +103,28 @@ def get_rope(
103 mrope_section=rope_parameters["mrope_section"],
104 mrope_interleaved=rope_parameters.get("mrope_interleaved", False),
105 )
106 + elif "use_fope" in rope_parameters and rope_parameters["use_fope"]:
107 + extra_kwargs = {
108 + k: v
109 + for k, v in rope_parameters.items()
110 + if k
111 + in (
112 + "num_key_value_heads",
113 + "num_inv_freq",
114 + "fope_sep_head",
115 + "fope_init_factor",
116 + )
117 + }
118 + extra_kwargs["init_cache"] = False
119 + rotary_emb = FourierRotaryEmbedding(
120 + head_size,
121 + rotary_dim,
122 + max_position,
123 + base,
124 + is_neox_style,
125 + dtype,
126 + **extra_kwargs,
127 + )
128 else:
129 rotary_emb = RotaryEmbedding(
130 head_size,
vllm/model_executor/layers/rotary_embedding/base.py
+15 -6
@@ -25,6 +25,7 @@ class RotaryEmbeddingBase(CustomOp):
25 base: float,
26 is_neox_style: bool,
27 dtype: torch.dtype,
28 + init_cache: bool = True,
29 ) -> None:
30 super().__init__()
31 self.head_size = head_size
@@ -46,11 +47,12 @@ class RotaryEmbeddingBase(CustomOp):
47 if not hasattr(self, "use_flashinfer"):
48 self.use_flashinfer = False
49
49 - cache = self._compute_cos_sin_cache()
50 - if not self.use_flashinfer:
51 - cache = cache.to(dtype)
52 - self.cos_sin_cache: torch.Tensor
53 - self.register_buffer("cos_sin_cache", cache, persistent=False)
50 + if init_cache:
51 + cache = self._compute_cos_sin_cache()
52 + if not self.use_flashinfer:
53 + cache = cache.to(dtype)
54 + self.cos_sin_cache: torch.Tensor
55 + self.register_buffer("cos_sin_cache", cache, persistent=False)
56 self.is_rocm_triton_rotary_embed_enabled = (
57 rocm_aiter_ops.is_triton_rotary_embed_enabled()
58 )
@@ -108,9 +110,16 @@ class RotaryEmbedding(RotaryEmbeddingBase):
110 base: float,
111 is_neox_style: bool,
112 dtype: torch.dtype,
113 + init_cache: bool = True,
114 ) -> None:
115 super().__init__(
113 - head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
116 + head_size=head_size,
117 + rotary_dim=rotary_dim,
118 + max_position_embeddings=max_position_embeddings,
119 + base=base,
120 + is_neox_style=is_neox_style,
121 + dtype=dtype,
122 + init_cache=init_cache,
123 )
124
125 @staticmethod
vllm/model_executor/layers/rotary_embedding/fope.py new
+199
@@ -0,0 +1,199 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +import torch.nn.functional as F
5 +from torch import nn
6 +
7 +from vllm.distributed import (
8 + get_tensor_model_parallel_rank,
9 + get_tensor_model_parallel_world_size,
10 +)
11 +
12 +from .base import RotaryEmbedding
13 +from .common import rotate_neox
14 +
15 +
16 +class FourierRotaryEmbedding(RotaryEmbedding):
17 + def __init__(
18 + self,
19 + head_size: int,
20 + rotary_dim: int,
21 + max_position_embeddings: int,
22 + base: float,
23 + is_neox_style: bool,
24 + dtype: torch.dtype,
25 + init_cache: bool,
26 + # extra parameters for FoPE
27 + num_key_value_heads: int,
28 + num_inv_freq: int,
29 + fope_sep_head: bool,
30 + fope_init_factor: float,
31 + ):
32 + # fope related parameters
33 + self.num_key_value_heads = num_key_value_heads
34 + self.num_inv_freq = num_inv_freq
35 + self.fope_sep_head = fope_sep_head
36 + self.fope_init_factor = fope_init_factor
37 +
38 + super().__init__(
39 + head_size=head_size,
40 + rotary_dim=rotary_dim,
41 + max_position_embeddings=max_position_embeddings,
42 + base=base,
43 + is_neox_style=is_neox_style,
44 + dtype=dtype,
45 + init_cache=init_cache,
46 + )
47 +
48 + # setup buffers and parameters
49 + self.inv_freq: torch.Tensor
50 + self.register_buffer(
51 + "inv_freq", self._compute_inv_freq(self.base), persistent=False
52 + )
53 +
54 + self.input_dim = self.inv_freq.shape[-1]
55 + self.output_dim = self.inv_freq.shape[-1]
56 + self.cos_coef = nn.Parameter(
57 + torch.empty(num_key_value_heads, self.input_dim, self.output_dim),
58 + requires_grad=False,
59 + )
60 + self.sin_coef = nn.Parameter(
61 + torch.empty(num_key_value_heads, self.input_dim, self.output_dim),
62 + requires_grad=False,
63 + )
64 + self.sin_coef.weight_loader = self.weight_loader
65 + self.cos_coef.weight_loader = self.weight_loader
66 +
67 + self.cos_sin_cache: torch.Tensor
68 + cache = self._compute_cos_sin_cache().to(dtype)
69 + self.register_buffer("cos_sin_cache", cache, persistent=False)
70 +
71 + # update cache in the first forward, where sin/cos_coef weights are ready
72 + self.update_cache = True
73 +
74 + def _compute_inv_freq(self, base: float) -> torch.Tensor:
75 + """Compute the inverse frequency."""
76 + inv_freq = 1.0 / (
77 + base
78 + ** (
79 + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
80 + )
81 + )
82 +
83 + inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool)
84 + if self.num_inv_freq is not None:
85 + inv_freq_idx_selected[self.num_inv_freq :] = False
86 + else:
87 + inv_freq_idx_selected = inv_freq > (
88 + 2.0 * torch.pi / self.max_position_embeddings
89 + )
90 +
91 + inv_freq = inv_freq[inv_freq_idx_selected]
92 + return inv_freq
93 +
94 + def _compute_cos_sin_cache(self) -> torch.Tensor:
95 + """Compute the cos and sin cache."""
96 + device = self.inv_freq.device
97 + t = torch.arange(self.max_position_embeddings, dtype=torch.float, device=device)
98 +
99 + freqs = torch.einsum("j,i -> ji", t, self.inv_freq)
100 + if self.fope_sep_head:
101 + pos_cos = freqs.cos().unsqueeze(0).expand(self.num_key_value_heads, -1, -1)
102 + pos_sin = freqs.sin().unsqueeze(0).expand(self.num_key_value_heads, -1, -1)
103 + else:
104 + pos_cos = freqs.cos()
105 + pos_sin = freqs.sin()
106 +
107 + if self.fope_sep_head:
108 + sin = torch.einsum("htD, hDd -> thd", pos_sin, self.sin_coef.float())
109 + cos = torch.einsum("htD, hDd -> thd", pos_cos, self.cos_coef.float())
110 + else:
111 + sin = torch.einsum("tD, Dd -> td", pos_sin, self.sin_coef.float())
112 + cos = torch.einsum("tD, Dd -> td", pos_cos, self.cos_coef.float())
113 +
114 + sin = F.pad(
115 + input=sin,
116 + pad=(0, self.head_size // 2 - sin.size(-1)),
117 + mode="constant",
118 + value=1,
119 + )
120 + cos = F.pad(
121 + input=cos,
122 + pad=(0, self.head_size // 2 - cos.size(-1)),
123 + mode="constant",
124 + value=1,
125 + )
126 +
127 + sin = torch.cat((sin, sin), dim=-1)
128 + cos = torch.cat((cos, cos), dim=-1)
129 +
130 + # cache: (max_position_embeddings, num_kv_heads, kv_size * 2)
131 + cache = torch.cat((cos, sin), dim=-1)
132 + return cache
133 +
134 + def forward_native(
135 + self,
136 + positions: torch.Tensor,
137 + query: torch.Tensor,
138 + key: torch.Tensor | None = None,
139 + offsets: torch.Tensor | None = None,
140 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
141 + # update cos/sin cache in the first forward
142 + if self.update_cache:
143 + cache = self._compute_cos_sin_cache().to(self.dtype)
144 + self.cos_sin_cache.copy_(cache)
145 + self.update_cache = False
146 +
147 + positions = positions.flatten()
148 + cos_sin = self.cos_sin_cache.index_select(0, positions)
149 + cos, sin = cos_sin.chunk(2, dim=-1)
150 +
151 + # apply rotary embedding
152 + # query: (seq_len, num_heads, head_size)
153 + # key: (seq_len, num_kv_heads, head_size)
154 + query = query.unflatten(-1, (-1, self.head_size))
155 + assert key is not None, "Key tensor is required for FoPE."
156 + key = key.unflatten(-1, (-1, self.head_size))
157 +
158 + assert query.dim() == key.dim() == 3, (
159 + "Expected query key (seq_len, heads, head_dim)"
160 + )
161 + assert cos.dim() <= 3 and sin.dim() <= 3
162 +
163 + need_reshape = False
164 + if cos.dim() == 3:
165 + # for fope
166 + need_reshape = True
167 + query_shape = query.shape
168 + key_shape = key.shape
169 + cos = cos.flatten(0, 1)
170 + sin = sin.flatten(0, 1)
171 + seq_len = cos.size(0)
172 + query = query.view(seq_len, -1, query.size(-1))
173 + key = key.view(seq_len, -1, key.size(-1))
174 +
175 + # native implementation of apply rope for neox style
176 + cos = cos.unsqueeze(1)
177 + sin = sin.unsqueeze(1)
178 + query = (query * cos) + (rotate_neox(query) * sin)
179 + key = (key * cos) + (rotate_neox(key) * sin)
180 +
181 + if need_reshape:
182 + query = query.view(query_shape)
183 + key = key.view(key_shape)
184 +
185 + return query, key
186 +
187 + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
188 + """load fope weights"""
189 + world_size = get_tensor_model_parallel_world_size()
190 + rank = get_tensor_model_parallel_rank()
191 + num_key_value_heads = loaded_weight.size(0)
192 +
193 + if num_key_value_heads < world_size:
194 + n_replicate = world_size // num_key_value_heads
195 + world_size = num_key_value_heads
196 + rank = rank // n_replicate
197 +
198 + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank]
199 + param.data.copy_(loaded_weight)
vllm/model_executor/models/interns1_pro.py new
+633
@@ -0,0 +1,633 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2025 The vLLM team.
5 +# Copyright 2025 The Qwen Team.
6 +# Copyright 2025 The HuggingFace Inc. team.
7 +# All rights reserved.
8 +#
9 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10 +# and OPT implementations in this library. It has been modified from its
11 +# original forms to accommodate minor architectural differences compared
12 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 +#
14 +# Licensed under the Apache License, Version 2.0 (the "License");
15 +# you may not use this file except in compliance with the License.
16 +# You may obtain a copy of the License at
17 +#
18 +# http://www.apache.org/licenses/LICENSE-2.0
19 +#
20 +# Unless required by applicable law or agreed to in writing, software
21 +# distributed under the License is distributed on an "AS IS" BASIS,
22 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 +# See the License for the specific language governing permissions and
24 +# limitations under the License.
25 +"""Inference-only InternS1Pro model compatible with HuggingFace weights."""
26 +
27 +import functools
28 +from collections.abc import Iterable
29 +from typing import Any
30 +
31 +import torch
32 +from torch import nn
33 +from transformers import AutoProcessor, PretrainedConfig
34 +
35 +from vllm.attention.layer import Attention
36 +from vllm.config import CacheConfig, VllmConfig
37 +from vllm.distributed import (
38 + get_ep_group,
39 + get_tensor_model_parallel_world_size,
40 + tensor_model_parallel_all_gather,
41 +)
42 +from vllm.logger import init_logger
43 +from vllm.model_executor.layers.activation import SiluAndMul
44 +from vllm.model_executor.layers.fused_moe import FusedMoE
45 +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
46 +from vllm.model_executor.layers.layernorm import RMSNorm
47 +from vllm.model_executor.layers.linear import (
48 + MergedColumnParallelLinear,
49 + QKVParallelLinear,
50 + ReplicatedLinear,
51 + RowParallelLinear,
52 +)
53 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
54 +from vllm.model_executor.layers.quantization import QuantizationConfig
55 +from vllm.model_executor.layers.rotary_embedding import get_rope
56 +from vllm.model_executor.layers.vocab_parallel_embedding import (
57 + ParallelLMHead,
58 +)
59 +from vllm.model_executor.models.utils import sequence_parallel_chunk
60 +from vllm.multimodal import MULTIMODAL_REGISTRY
61 +
62 +from .interfaces import MixtureOfExperts
63 +from .qwen3_moe import (
64 + Qwen3MoeForCausalLM,
65 +)
66 +from .qwen3_vl import (
67 + Qwen3_VisionTransformer,
68 + Qwen3VLDummyInputsBuilder,
69 + Qwen3VLForConditionalGeneration,
70 + Qwen3VLMultiModalProcessor,
71 + Qwen3VLProcessingInfo,
72 +)
73 +from .qwen3_vl_moe import Qwen3MoeLLMModel
74 +from .utils import (
75 + AutoWeightsLoader,
76 + WeightsMapper,
77 + extract_layer_index,
78 + maybe_prefix,
79 +)
80 +
81 +logger = init_logger(__name__)
82 +
83 +
84 +class InternS1ProProcessingInfo(Qwen3VLProcessingInfo):
85 + def get_hf_config(self):
86 + return self.ctx.get_hf_config()
87 +
88 + def get_hf_processor(self, **kwargs: object) -> AutoProcessor:
89 + return AutoProcessor.from_pretrained(
90 + self.ctx.model_config.model,
91 + trust_remote_code=True,
92 + **kwargs,
93 + )
94 +
95 +
96 +class InternS1ProMoeMLP(nn.Module):
97 + def __init__(
98 + self,
99 + hidden_size: int,
100 + intermediate_size: int,
101 + hidden_act: str,
102 + quant_config: QuantizationConfig | None = None,
103 + reduce_results: bool = True,
104 + prefix: str = "",
105 + ) -> None:
106 + super().__init__()
107 + self.gate_up_proj = MergedColumnParallelLinear(
108 + hidden_size,
109 + [intermediate_size] * 2,
110 + bias=False,
111 + quant_config=quant_config,
112 + prefix=f"{prefix}.gate_up_proj",
113 + )
114 + self.down_proj = RowParallelLinear(
115 + intermediate_size,
116 + hidden_size,
117 + bias=False,
118 + quant_config=quant_config,
119 + reduce_results=reduce_results,
120 + prefix=f"{prefix}.down_proj",
121 + )
122 + if hidden_act != "silu":
123 + raise ValueError(
124 + f"Unsupported activation: {hidden_act}. Only silu is supported for now."
125 + )
126 + self.act_fn = SiluAndMul()
127 +
128 + def forward(self, x):
129 + gate_up, _ = self.gate_up_proj(x)
130 + x = self.act_fn(gate_up)
131 + x, _ = self.down_proj(x)
132 + return x
133 +
134 +
135 +class InternS1ProMoeSparseMoeBlock(nn.Module):
136 + def __init__(
137 + self,
138 + vllm_config: VllmConfig,
139 + prefix: str = "",
140 + ):
141 + super().__init__()
142 +
143 + config = vllm_config.model_config.hf_text_config
144 + parallel_config = vllm_config.parallel_config
145 + quant_config = vllm_config.quant_config
146 +
147 + self.tp_size = get_tensor_model_parallel_world_size()
148 +
149 + self.ep_group = get_ep_group().device_group
150 + self.ep_rank = get_ep_group().rank_in_group
151 + self.ep_size = self.ep_group.size()
152 + self.n_routed_experts = config.num_experts
153 +
154 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
155 +
156 + if self.tp_size > config.num_experts:
157 + raise ValueError(
158 + f"Tensor parallel size {self.tp_size} is greater than "
159 + f"the number of experts {config.num_experts}."
160 + )
161 +
162 + # Load balancing settings.
163 + eplb_config = vllm_config.parallel_config.eplb_config
164 + self.enable_eplb = parallel_config.enable_eplb
165 +
166 + self.n_logical_experts = self.n_routed_experts
167 + self.n_redundant_experts = eplb_config.num_redundant_experts
168 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
169 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
170 +
171 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
172 + self.physical_expert_end = (
173 + self.physical_expert_start + self.n_local_physical_experts
174 + )
175 +
176 + # For custom routing function
177 + self.n_groups = getattr(config, "router_n_groups", -1)
178 +
179 + self.experts = FusedMoE(
180 + num_experts=self.n_routed_experts,
181 + top_k=config.num_experts_per_tok,
182 + hidden_size=config.hidden_size,
183 + intermediate_size=config.moe_intermediate_size,
184 + reduce_results=True,
185 + renormalize=config.norm_topk_prob,
186 + quant_config=quant_config,
187 + prefix=f"{prefix}.experts",
188 + enable_eplb=self.enable_eplb,
189 + num_redundant_experts=self.n_redundant_experts,
190 + is_sequence_parallel=self.is_sequence_parallel,
191 + routing_method_type=RoutingMethodType.Renormalize,
192 + custom_routing_function=self._custom_routing_function,
193 + )
194 +
195 + self.gate = ReplicatedLinear(
196 + config.hidden_size,
197 + config.num_experts,
198 + bias=False,
199 + prefix=f"{prefix}.gate",
200 + )
201 +
202 + @staticmethod
203 + @functools.lru_cache
204 + def get_group_offsets(n_groups: int, group_size: int, device: str):
205 + group_offsets = (torch.arange(n_groups, device=device) * group_size).view(
206 + 1, -1, 1
207 + ) # [1, n_groups, 1]
208 + return group_offsets
209 +
210 + # TODO: zhouxinyu, use vllm routing functions
211 + def _custom_routing_function(
212 + self,
213 + hidden_states: torch.Tensor,
214 + gating_output: torch.Tensor,
215 + topk: int,
216 + renormalize: bool,
217 + ) -> torch.Tensor:
218 + routing_weights = torch.softmax(gating_output, dim=-1, dtype=torch.float32)
219 +
220 + if self.n_groups > 0:
221 + assert routing_weights.shape[-1] % self.n_groups == 0, (
222 + f"{routing_weights.shape[-1]} cannot be divided by {self.n_groups}"
223 + )
224 + per_group_top_k = topk // self.n_groups
225 + group_size = routing_weights.shape[-1] // self.n_groups
226 + group_offsets = self.get_group_offsets(
227 + self.n_groups, group_size, routing_weights.device
228 + )
229 + routing_weights = routing_weights.unflatten(-1, (self.n_groups, group_size))
230 + topk_weights, topk_ids = torch.topk(
231 + routing_weights, per_group_top_k, dim=-1
232 + )
233 + topk_ids = (topk_ids + group_offsets).flatten(-2, -1)
234 + topk_weights = topk_weights.flatten(-2, -1)
235 + else:
236 + topk_weights, topk_ids = torch.topk(routing_weights, topk, dim=-1)
237 +
238 + if renormalize:
239 + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
240 +
241 + return topk_weights, topk_ids
242 +
243 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
244 + assert hidden_states.dim() <= 2, (
245 + "InternS1ProMoeSparseMoeBlock only supports 1D or 2D inputs"
246 + )
247 + is_input_1d = hidden_states.dim() == 1
248 + num_tokens, hidden_dim = hidden_states.shape
249 + hidden_states = hidden_states.view(-1, hidden_dim)
250 +
251 + if self.is_sequence_parallel:
252 + hidden_states = sequence_parallel_chunk(hidden_states)
253 +
254 + # router_logits: (num_tokens, n_experts)
255 + router_logits, _ = self.gate(hidden_states)
256 + final_hidden_states = self.experts(
257 + hidden_states=hidden_states, router_logits=router_logits
258 + )
259 +
260 + if self.is_sequence_parallel:
261 + final_hidden_states = tensor_model_parallel_all_gather(
262 + final_hidden_states, 0
263 + )
264 + final_hidden_states = final_hidden_states[:num_tokens]
265 +
266 + # return to 1d if input is 1d
267 + return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
268 +
269 +
270 +class InternS1ProMoeAttention(nn.Module):
271 + def __init__(
272 + self,
273 + hidden_size: int,
274 + num_heads: int,
275 + num_kv_heads: int,
276 + rope_parameters: dict[str, Any],
277 + max_position_embeddings: int = 32768,
278 + head_dim: int | None = None,
279 + rms_norm_eps: float = 1e-06,
280 + qkv_bias: bool = False,
281 + cache_config: CacheConfig | None = None,
282 + quant_config: QuantizationConfig | None = None,
283 + prefix: str = "",
284 + dual_chunk_attention_config: dict[str, Any] | None = None,
285 + ) -> None:
286 + super().__init__()
287 + self.hidden_size = hidden_size
288 + tp_size = get_tensor_model_parallel_world_size()
289 + self.total_num_heads = num_heads
290 + assert self.total_num_heads % tp_size == 0
291 + self.num_heads = self.total_num_heads // tp_size
292 + self.total_num_kv_heads = num_kv_heads
293 + if self.total_num_kv_heads >= tp_size:
294 + # Number of KV heads is greater than TP size, so we partition
295 + # the KV heads across multiple tensor parallel GPUs.
296 + assert self.total_num_kv_heads % tp_size == 0
297 + else:
298 + # Number of KV heads is less than TP size, so we replicate
299 + # the KV heads across multiple tensor parallel GPUs.
300 + assert tp_size % self.total_num_kv_heads == 0
301 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
302 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)
303 + self.q_size = self.num_heads * self.head_dim
304 + self.kv_size = self.num_kv_heads * self.head_dim
305 + self.scaling = self.head_dim**-0.5
306 + self.max_position_embeddings = max_position_embeddings
307 + self.dual_chunk_attention_config = dual_chunk_attention_config
308 +
309 + self.qkv_proj = QKVParallelLinear(
310 + hidden_size,
311 + self.head_dim,
312 + self.total_num_heads,
313 + self.total_num_kv_heads,
314 + bias=qkv_bias,
315 + quant_config=quant_config,
316 + prefix=f"{prefix}.qkv_proj",
317 + )
318 +
319 + self.o_proj = RowParallelLinear(
320 + self.total_num_heads * self.head_dim,
321 + hidden_size,
322 + bias=False,
323 + quant_config=quant_config,
324 + prefix=f"{prefix}.o_proj",
325 + )
326 +
327 + rope_parameters["num_key_value_heads"] = self.num_kv_heads
328 + self.rotary_emb = get_rope(
329 + self.head_dim,
330 + max_position=max_position_embeddings,
331 + rope_parameters=rope_parameters,
332 + dual_chunk_attention_config=dual_chunk_attention_config,
333 + )
334 +
335 + self.attn = Attention(
336 + self.num_heads,
337 + self.head_dim,
338 + self.scaling,
339 + num_kv_heads=self.num_kv_heads,
340 + cache_config=cache_config,
341 + quant_config=quant_config,
342 + prefix=f"{prefix}.attn",
343 + **{
344 + "layer_idx": extract_layer_index(prefix),
345 + "dual_chunk_attention_config": dual_chunk_attention_config,
346 + }
347 + if dual_chunk_attention_config
348 + else {},
349 + )
350 +
351 + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
352 + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
353 +
354 + def forward(
355 + self,
356 + positions: torch.Tensor,
357 + hidden_states: torch.Tensor,
358 + ) -> torch.Tensor:
359 + qkv, _ = self.qkv_proj(hidden_states)
360 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
361 + # Add qk-norm
362 + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
363 + q_by_head = self.q_norm(q_by_head)
364 + q = q_by_head.view(q.shape)
365 +
366 + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)
367 + k_by_head = self.k_norm(k_by_head)
368 + k = k_by_head.view(k.shape)
369 + q, k = self.rotary_emb.forward_native(positions, q, k)
370 + attn_output = self.attn(q, k, v)
371 + output, _ = self.o_proj(attn_output)
372 + return output
373 +
374 +
375 +class InternS1ProMoeDecoderLayer(nn.Module):
376 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
377 + super().__init__()
378 +
379 + config = vllm_config.model_config.hf_text_config
380 + cache_config = vllm_config.cache_config
381 + quant_config = vllm_config.quant_config
382 +
383 + self.hidden_size = config.hidden_size
384 + max_position_embeddings = getattr(config, "max_position_embeddings", 32768)
385 + dual_chunk_attention_config = getattr(
386 + config, "dual_chunk_attention_config", None
387 + )
388 +
389 + # update rope related parameters
390 + rope_scaling = config.rope_scaling
391 + fope_keys = {"fope_init_factor", "fope_sep_head", "num_inv_freq"}
392 + use_fope = any(rope_scaling.get(key) is not None for key in fope_keys)
393 + fope_init_factor = rope_scaling.get("fope_init_factor", None)
394 + fope_sep_head = rope_scaling.get("fope_sep_head", None)
395 + num_inv_freq = rope_scaling.get("num_inv_freq", None)
396 +
397 + config.rope_parameters["use_fope"] = use_fope
398 + config.rope_parameters["fope_init_factor"] = fope_init_factor
399 + config.rope_parameters["fope_sep_head"] = fope_sep_head
400 + config.rope_parameters["num_inv_freq"] = num_inv_freq
401 +
402 + assert use_fope, "should use FOPE for InternS1Pro model"
403 + self.self_attn = InternS1ProMoeAttention(
404 + hidden_size=self.hidden_size,
405 + num_heads=config.num_attention_heads,
406 + num_kv_heads=config.num_key_value_heads,
407 + rope_parameters=config.rope_parameters,
408 + max_position_embeddings=max_position_embeddings,
409 + rms_norm_eps=config.rms_norm_eps,
410 + qkv_bias=getattr(config, "attention_bias", False),
411 + head_dim=getattr(config, "head_dim", None),
412 + cache_config=cache_config,
413 + quant_config=quant_config,
414 + prefix=f"{prefix}.self_attn",
415 + dual_chunk_attention_config=dual_chunk_attention_config,
416 + )
417 +
418 + # `mlp_only_layers` in the config.
419 + layer_idx = extract_layer_index(prefix)
420 + mlp_only_layers = (
421 + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
422 + )
423 + if (layer_idx not in mlp_only_layers) and (
424 + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
425 + ):
426 + self.mlp = InternS1ProMoeSparseMoeBlock(
427 + vllm_config=vllm_config, prefix=f"{prefix}.mlp"
428 + )
429 + else:
430 + self.mlp = InternS1ProMoeMLP(
431 + hidden_size=config.hidden_size,
432 + intermediate_size=config.intermediate_size,
433 + hidden_act=config.hidden_act,
434 + quant_config=quant_config,
435 + prefix=f"{prefix}.mlp",
436 + )
437 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
438 + self.post_attention_layernorm = RMSNorm(
439 + config.hidden_size, eps=config.rms_norm_eps
440 + )
441 +
442 + def forward(
443 + self,
444 + positions: torch.Tensor,
445 + hidden_states: torch.Tensor,
446 + residual: torch.Tensor | None,
447 + ) -> tuple[torch.Tensor, torch.Tensor]:
448 + # Self Attention
449 + if residual is None:
450 + residual = hidden_states
451 + hidden_states = self.input_layernorm(hidden_states)
452 + else:
453 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
454 + hidden_states = self.self_attn(
455 + positions=positions,
456 + hidden_states=hidden_states,
457 + )
458 +
459 + # Fully Connected
460 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
461 + hidden_states = self.mlp(hidden_states)
462 + return hidden_states, residual
463 +
464 +
465 +class InternS1ProMoeLLMModel(Qwen3MoeLLMModel):
466 + def __init__(
467 + self,
468 + *,
469 + vllm_config: VllmConfig,
470 + prefix: str = "",
471 + decoder_layer_type: type[torch.nn.Module] = InternS1ProMoeDecoderLayer,
472 + ):
473 + super().__init__(
474 + vllm_config=vllm_config,
475 + prefix=prefix,
476 + decoder_layer_type=decoder_layer_type,
477 + )
478 +
479 +
480 +class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM):
481 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
482 + super().__init__()
483 + self.config = vllm_config.model_config.hf_config.text_config
484 + self.quant_config = vllm_config.quant_config
485 + self.model = InternS1ProMoeLLMModel(
486 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
487 + )
488 + self.lm_head = ParallelLMHead(
489 + self.config.vocab_size,
490 + self.config.hidden_size,
491 + quant_config=self.quant_config,
492 + prefix=maybe_prefix(prefix, "lm_head"),
493 + )
494 + if self.config.tie_word_embeddings:
495 + self.lm_head.weight = self.model.embed_tokens.weight
496 + self.logits_processor = LogitsProcessor(self.config.vocab_size)
497 + self.make_empty_intermediate_tensors = (
498 + self.model.make_empty_intermediate_tensors
499 + )
500 +
501 +
502 +class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts):
503 + def update_physical_experts_metadata(
504 + self,
505 + num_physical_experts: int,
506 + num_local_physical_experts: int,
507 + ) -> None:
508 + assert self.num_local_physical_experts == num_local_physical_experts
509 + self.num_physical_experts = num_physical_experts
510 + self.num_local_physical_experts = num_local_physical_experts
511 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
512 + for layer in self.language_model.model.layers:
513 + if isinstance(layer.mlp, InternS1ProMoeSparseMoeBlock):
514 + moe = layer.mlp
515 + moe.n_local_physical_experts = num_local_physical_experts
516 + moe.n_physical_experts = num_physical_experts
517 + moe.n_redundant_experts = self.num_redundant_experts
518 + moe.experts.update_expert_map()
519 +
520 + def set_moe_parameters(self):
521 + self.expert_weights = []
522 +
523 + self.moe_layers = []
524 + example_moe = None
525 + for layer in self.language_model.model.layers:
526 + if hasattr(layer, "mlp") and isinstance(
527 + layer.mlp, InternS1ProMoeSparseMoeBlock
528 + ):
529 + example_moe = layer.mlp
530 + self.moe_layers.append(layer.mlp.experts)
531 +
532 + if example_moe is None:
533 + raise RuntimeError("No InternS1ProMoe layer found in the language_model.")
534 +
535 + # Set MoE hyperparameters
536 + self.num_moe_layers = len(self.moe_layers)
537 + self.num_expert_groups = 1
538 + self.num_shared_experts = 0
539 + self.num_logical_experts = example_moe.n_logical_experts
540 + self.num_physical_experts = example_moe.n_physical_experts
541 + self.num_local_physical_experts = example_moe.n_local_physical_experts
542 + self.num_routed_experts = example_moe.n_routed_experts
543 + self.num_redundant_experts = example_moe.n_redundant_experts
544 +
545 +
546 +@MULTIMODAL_REGISTRY.register_processor(
547 + Qwen3VLMultiModalProcessor,
548 + info=InternS1ProProcessingInfo,
549 + dummy_inputs=Qwen3VLDummyInputsBuilder,
550 +)
551 +class InternS1ProForConditionalGeneration(
552 + Qwen3VLForConditionalGeneration, Qwen3VLMoeMixtureOfExperts
553 +):
554 + is_3d_moe_weight: bool = True
555 + packed_modules_mapping = {
556 + "qkv_proj": [
557 + "q_proj",
558 + "k_proj",
559 + "v_proj",
560 + ],
561 + }
562 +
563 + # To ensure correct weight loading and mapping.
564 + hf_to_vllm_mapper = WeightsMapper(
565 + orig_to_new_prefix={
566 + "model.visual.": "visual.",
567 + "lm_head.": "language_model.lm_head.",
568 + "model.language_model.": "language_model.model.",
569 + },
570 + orig_to_new_suffix={
571 + # Handle FOPE rotary embeddings
572 + ".rotary_emb.sin_coef": ".layers.0.self_attn.rotary_emb.sin_coef",
573 + ".rotary_emb.cos_coef": ".layers.0.self_attn.rotary_emb.cos_coef",
574 + },
575 + )
576 +
577 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
578 + super().__init__()
579 + config: PretrainedConfig = vllm_config.model_config.hf_config
580 + multimodal_config = vllm_config.model_config.multimodal_config
581 +
582 + self.config = config
583 + self.multimodal_config = multimodal_config
584 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
585 + self.video_pruning_rate = multimodal_config.video_pruning_rate
586 + self.is_multimodal_pruning_enabled = (
587 + multimodal_config.is_multimodal_pruning_enabled()
588 + )
589 +
590 + if not multimodal_config.get_limit_per_prompt(
591 + "image"
592 + ) and not multimodal_config.get_limit_per_prompt("video"):
593 + self.visual = None
594 + else:
595 + self.visual = Qwen3_VisionTransformer(
596 + config.vision_config,
597 + norm_eps=getattr(config, "rms_norm_eps", 1e-6),
598 + multimodal_config=multimodal_config,
599 + prefix=maybe_prefix(prefix, "visual"),
600 + )
601 +
602 + self.language_model = InternS1ProMoeLLMForCausalLM(
603 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
604 + )
605 + # Whether to include the gate_up_proj mapping is determined by
606 + # the language model.
607 + self.packed_modules_mapping = (
608 + self.packed_modules_mapping | self.language_model.packed_modules_mapping
609 + )
610 +
611 + self.make_empty_intermediate_tensors = (
612 + self.language_model.make_empty_intermediate_tensors
613 + )
614 +
615 + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes")
616 + self.deepstack_num_level = (
617 + len(config.vision_config.deepstack_visual_indexes)
618 + if self.use_deepstack
619 + else 0
620 + )
621 + self.visual_dim = config.vision_config.out_hidden_size
622 + self.multiscale_dim = self.visual_dim * self.deepstack_num_level
623 +
624 + # Set MoE hyperparameters
625 + self.set_moe_parameters()
626 +
627 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
628 + """load weights"""
629 + skip_prefixes = ["model.time_series."]
630 + if self.visual is None:
631 + skip_prefixes.append("visual.")
632 + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
633 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/qwen3_moe.py
+8 -2
@@ -428,7 +428,13 @@ class Qwen3MoeDecoderLayer(nn.Module):
428
429 @support_torch_compile
430 class Qwen3MoeModel(nn.Module):
431 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
431 + def __init__(
432 + self,
433 + *,
434 + vllm_config: VllmConfig,
435 + prefix: str = "",
436 + decoder_layer_type: type[torch.nn.Module] = Qwen3MoeDecoderLayer,
437 + ):
438 super().__init__()
439
440 config = vllm_config.model_config.hf_text_config
@@ -449,7 +455,7 @@ class Qwen3MoeModel(nn.Module):
455 )
456 self.start_layer, self.end_layer, self.layers = make_layers(
457 config.num_hidden_layers,
452 - lambda prefix: Qwen3MoeDecoderLayer(vllm_config=vllm_config, prefix=prefix),
458 + lambda prefix: decoder_layer_type(vllm_config=vllm_config, prefix=prefix),
459 prefix=f"{prefix}.layers",
460 )
461 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
vllm/model_executor/models/qwen3_vl.py
+5 -1
@@ -325,7 +325,11 @@ class Qwen3_VisionTransformer(nn.Module):
325 self.spatial_merge_size = vision_config.spatial_merge_size
326 self.spatial_merge_unit = self.spatial_merge_size**2
327 self.temporal_patch_size = vision_config.temporal_patch_size
328 - self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes
328 + self.deepstack_visual_indexes = (
329 + vision_config.deepstack_visual_indexes
330 + if hasattr(vision_config, "deepstack_visual_indexes")
331 + else []
332 + )
333 self.num_grid_per_side = int(self.num_position_embeddings**0.5)
334
335 # NOTE: This is used for creating empty tensor for all_gather for
vllm/model_executor/models/qwen3_vl_moe.py
+13 -2
@@ -48,6 +48,7 @@ from vllm.sequence import IntermediateTensors
48
49 from .interfaces import MixtureOfExperts
50 from .qwen3_moe import (
51 + Qwen3MoeDecoderLayer,
52 Qwen3MoeForCausalLM,
53 Qwen3MoeModel,
54 Qwen3MoeSparseMoeBlock,
@@ -82,8 +83,18 @@ class Qwen3VLMoeProcessingInfo(Qwen3VLProcessingInfo):
83 }
84 )
85 class Qwen3MoeLLMModel(Qwen3MoeModel):
85 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
86 - super().__init__(vllm_config=vllm_config, prefix=prefix)
86 + def __init__(
87 + self,
88 + *,
89 + vllm_config: VllmConfig,
90 + prefix: str = "",
91 + decoder_layer_type: type[torch.nn.Module] = Qwen3MoeDecoderLayer,
92 + ):
93 + super().__init__(
94 + vllm_config=vllm_config,
95 + prefix=prefix,
96 + decoder_layer_type=decoder_layer_type,
97 + )
98 if not get_pp_group().is_first_rank:
99 assert self.start_layer >= len(
100 vllm_config.model_config.hf_config.vision_config.deepstack_visual_indexes
vllm/model_executor/models/registry.py
+4
@@ -357,6 +357,10 @@ _MULTIMODAL_MODELS = {
357 "interns1",
358 "InternS1ForConditionalGeneration",
359 ),
360 + "InternS1ProForConditionalGeneration": (
361 + "interns1_pro",
362 + "InternS1ProForConditionalGeneration",
363 + ),
364 "Idefics3ForConditionalGeneration": (
365 "idefics3",
366 "Idefics3ForConditionalGeneration",