Interns2mobius support (#51149)
Signed-off-by: Lyu Han <lvhan_028@163.com>
Lyu Han committed
Aug 6, 2026 at 18:39 UTC
22013f74ffb0ae22deb29233cbe9fcd3efb1b374
8 files changed
+664
-3
docs/models/supported_models.md
+1
@@ -540,6 +540,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
540
| `IsaacForConditionalGeneration` | Isaac | T + I<sup>+</sup> | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ |
541
| `InternS1ForConditionalGeneration` | Intern-S1 | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |
542
| `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ |
543
+| `InternS2MobiusForConditionalGeneration` | Intern-S2-Mobius | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S2-Mobius` | ✅︎ | ✅︎ |
544
| `InternS2PreviewForConditionalGeneration` | Intern-S2-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S2-Preview`, etc. | ✅︎ | ✅︎ |
545
| `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, 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/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
546
| `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+11
@@ -996,6 +996,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
996
"internlm/Intern-S1-Pro",
997
trust_remote_code=True,
998
),
999
+ "InternS2MobiusForConditionalGeneration": _HfExamplesInfo(
1000
+ "internlm/Intern-S2-Mobius",
1001
+ trust_remote_code=True,
1002
+ is_available_online=False,
1003
+ ),
1004
"InternS2PreviewForConditionalGeneration": _HfExamplesInfo(
1005
"internlm/Intern-S2-Preview",
1006
trust_remote_code=True,
@@ -1682,6 +1687,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1687
trust_remote_code=True,
1688
max_model_len=4096,
1689
),
1690
+ "InternS2MobiusMTP": _HfExamplesInfo(
1691
+ "internlm/Intern-S2-Mobius",
1692
+ trust_remote_code=True,
1693
+ speculative_model="internlm/Intern-S2-Mobius",
1694
+ is_available_online=False,
1695
+ ),
1696
"KimiK3MTPModel": _HfExamplesInfo(
1697
"moonshotai/Kimi-K3",
1698
speculative_model="moonshotai/Kimi-K3",
vllm/config/speculative.py
+16
-3
@@ -538,15 +538,28 @@ class SpeculativeConfig:
538
"architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],
539
}
540
)
541
- if hf_config.model_type == "intern_s2_preview":
541
+ if hf_config.model_type in ("intern_s2_preview", "interns2_mobius"):
542
+ is_mobius = hf_config.model_type == "interns2_mobius"
543
text_config = getattr(hf_config, "text_config", None)
543
- is_moe = getattr(text_config, "model_type", None) == "qwen3_5_moe_text"
544
+ is_moe = is_mobius or (
545
+ getattr(text_config, "model_type", None) == "qwen3_5_moe_text"
546
+ )
547
+ if is_mobius:
548
+ assert text_config is not None
549
+ text_config.model_type = "qwen3_5_moe_text"
550
+ text_config.num_experts = text_config.mtp_num_experts
551
+ text_config.num_experts_per_tok = text_config.mtp_num_experts_per_tok
552
hf_config.model_type = "qwen3_5_mtp"
553
n_predict = getattr(text_config, "mtp_num_hidden_layers", None)
554
+ architecture = (
555
+ "InternS2MobiusMTP"
556
+ if is_mobius
557
+ else ("Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP")
558
+ )
559
hf_config.update(
560
{
561
"n_predict": n_predict,
549
- "architectures": ["Qwen3_5MoeMTP" if is_moe else "Qwen3_5MTP"],
562
+ "architectures": [architecture],
563
}
564
)
565
if hf_config.model_type in ("longcat_flash", "longcat_flash_ngram"):
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+25
@@ -62,6 +62,31 @@ from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
62
class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
63
"""Triton-based fused MoE expert implementation."""
64
65
+ @staticmethod
66
+ def is_supported_config(
67
+ cls: type[mk.FusedMoEExperts],
68
+ moe_config: FusedMoEConfig,
69
+ weight_key: QuantKey | None,
70
+ activation_key: QuantKey | None,
71
+ activation_format: mk.FusedMoEActivationFormat,
72
+ ) -> tuple[bool, str | None]:
73
+ supported, reason = mk.FusedMoEExperts.is_supported_config(
74
+ cls,
75
+ moe_config,
76
+ weight_key,
77
+ activation_key,
78
+ activation_format,
79
+ )
80
+ if not supported or not current_platform.is_cuda_alike():
81
+ return supported, reason
82
+
83
+ padded_num_experts = (moe_config.num_experts + 31) // 32 * 32
84
+ if padded_num_experts >= 1024:
85
+ return False, (
86
+ "kernel's moe_align_block_size requires fewer than 1024 padded experts"
87
+ )
88
+ return True, None
89
+
90
def __init__(
91
self,
92
moe_config: FusedMoEConfig,
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+22
@@ -34,6 +34,28 @@ class TrtLlmBf16ExpertsBase:
34
monolithic interfaces.
35
"""
36
37
+ @staticmethod
38
+ def is_supported_config(
39
+ cls: type[mk.FusedMoEExperts],
40
+ moe_config: FusedMoEConfig,
41
+ weight_key: QuantKey | None,
42
+ activation_key: QuantKey | None,
43
+ activation_format: mk.FusedMoEActivationFormat,
44
+ ) -> tuple[bool, str | None]:
45
+ supported, reason = mk.FusedMoEExperts.is_supported_config(
46
+ cls,
47
+ moe_config,
48
+ weight_key,
49
+ activation_key,
50
+ activation_format,
51
+ )
52
+ if not supported or moe_config.num_experts <= 2048:
53
+ return supported, reason
54
+ return False, (
55
+ "FlashInfer TRTLLM routing supports at most 2048 experts, "
56
+ f"but got {moe_config.num_experts}"
57
+ )
58
+
59
def supports_routing_replay_capture(self) -> bool:
60
return True
61
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+22
@@ -42,6 +42,28 @@ class TrtLlmFp8ExpertsBase:
42
interfaces.
43
"""
44
45
+ @staticmethod
46
+ def is_supported_config(
47
+ cls: type[mk.FusedMoEExperts],
48
+ moe_config: FusedMoEConfig,
49
+ weight_key: QuantKey | None,
50
+ activation_key: QuantKey | None,
51
+ activation_format: mk.FusedMoEActivationFormat,
52
+ ) -> tuple[bool, str | None]:
53
+ supported, reason = mk.FusedMoEExperts.is_supported_config(
54
+ cls,
55
+ moe_config,
56
+ weight_key,
57
+ activation_key,
58
+ activation_format,
59
+ )
60
+ if not supported or moe_config.num_experts <= 2048:
61
+ return supported, reason
62
+ return False, (
63
+ "FlashInfer TRTLLM routing supports at most 2048 experts, "
64
+ f"but got {moe_config.num_experts}"
65
+ )
66
+
67
def __init__(
68
self,
69
moe_config: FusedMoEConfig,
vllm/model_executor/models/interns2_mobius.py
new
+562
@@ -0,0 +1,562 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only Intern-S2-Mobius model."""
4
+
5
+from collections.abc import Iterable
6
+from itertools import islice
7
+
8
+import torch
9
+from torch import nn
10
+
11
+from vllm.compilation.decorators import support_torch_compile
12
+from vllm.config import VllmConfig, replace, set_current_vllm_config
13
+from vllm.distributed import (
14
+ get_ep_group,
15
+ get_pp_group,
16
+ get_tensor_model_parallel_world_size,
17
+ tensor_model_parallel_all_gather,
18
+ tensor_model_parallel_reduce_scatter,
19
+)
20
+from vllm.model_executor.layers.fused_moe import FusedMoEFactory
21
+from vllm.model_executor.layers.linear import ReplicatedLinear
22
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
23
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
24
+ QwenGatedDeltaNetAttention,
25
+)
26
+from vllm.model_executor.layers.vocab_parallel_embedding import (
27
+ ParallelLMHead,
28
+ VocabParallelEmbedding,
29
+)
30
+from vllm.multimodal import MULTIMODAL_REGISTRY
31
+from vllm.sequence import IntermediateTensors
32
+from vllm.tokenizers.registry import cached_tokenizer_from_config
33
+
34
+from .interfaces import EagleModelMixin
35
+from .interns2_preview import (
36
+ InternS2PreviewForConditionalGeneration,
37
+ InternS2PreviewProcessingInfo,
38
+)
39
+from .qwen2_moe import Qwen2MoeMLP
40
+from .qwen3_5 import (
41
+ Qwen3_5ForCausalLMBase,
42
+ Qwen3_5Model,
43
+ Qwen3_5RMSNorm,
44
+)
45
+from .qwen3_5_mtp import Qwen3_5MoeMTP
46
+from .qwen3_next import Qwen3NextAttention, _all_gather_hidden_and_residual
47
+from .qwen3_vl import (
48
+ Qwen3_VisionTransformer,
49
+ Qwen3VLDummyInputsBuilder,
50
+ Qwen3VLMultiModalProcessor,
51
+)
52
+from .utils import (
53
+ AutoWeightsLoader,
54
+ PPMissingLayer,
55
+ extract_layer_index,
56
+ make_empty_intermediate_tensors_factory,
57
+ make_layers,
58
+ maybe_prefix,
59
+ sequence_parallel_chunk,
60
+)
61
+
62
+
63
+class InternS2MobiusProcessingInfo(InternS2PreviewProcessingInfo):
64
+ pass
65
+
66
+
67
+class InternS2MobiusMetaMoeBlock(nn.Module):
68
+ """A routed MoE bank shared by multiple decoder layers."""
69
+
70
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
71
+ super().__init__()
72
+
73
+ config = vllm_config.model_config.hf_text_config
74
+ parallel_config = vllm_config.parallel_config
75
+
76
+ self.tp_size = get_tensor_model_parallel_world_size()
77
+ self.ep_group = get_ep_group().device_group
78
+ self.ep_rank = get_ep_group().rank_in_group
79
+ self.ep_size = self.ep_group.size()
80
+ self.n_routed_experts = config.num_experts
81
+ self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
82
+
83
+ if self.tp_size > self.n_routed_experts:
84
+ raise ValueError(
85
+ f"Tensor parallel size {self.tp_size} is greater than "
86
+ f"the number of experts {self.n_routed_experts}."
87
+ )
88
+
89
+ eplb_config = parallel_config.eplb_config
90
+ self.enable_eplb = parallel_config.enable_eplb
91
+ self.n_logical_experts = self.n_routed_experts
92
+ self.n_redundant_experts = eplb_config.num_redundant_experts
93
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
94
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
95
+
96
+ self.gate = ReplicatedLinear(
97
+ config.hidden_size,
98
+ config.num_experts,
99
+ bias=False,
100
+ quant_config=None,
101
+ prefix=f"{prefix}.gate",
102
+ )
103
+ self.experts = FusedMoEFactory(
104
+ gate=self.gate,
105
+ num_experts=config.num_experts,
106
+ top_k=config.num_experts_per_tok,
107
+ hidden_size=config.hidden_size,
108
+ intermediate_size=config.moe_intermediate_size,
109
+ renormalize=getattr(config, "norm_topk_prob", True),
110
+ quant_config=vllm_config.quant_config,
111
+ prefix=f"{prefix}.experts",
112
+ enable_eplb=self.enable_eplb,
113
+ num_redundant_experts=self.n_redundant_experts,
114
+ is_sequence_parallel=self.is_sequence_parallel,
115
+ activation=config.hidden_act,
116
+ )
117
+
118
+ def forward(
119
+ self,
120
+ hidden_states: torch.Tensor,
121
+ already_sequence_parallel: bool = False,
122
+ ) -> torch.Tensor:
123
+ orig_shape = hidden_states.shape
124
+ num_tokens, hidden_dim = hidden_states.shape
125
+ hidden_states = hidden_states.view(-1, hidden_dim)
126
+
127
+ if self.is_sequence_parallel and not already_sequence_parallel:
128
+ hidden_states = sequence_parallel_chunk(hidden_states)
129
+
130
+ if self.experts.is_internal_router:
131
+ final_hidden_states = self.experts(
132
+ hidden_states=hidden_states,
133
+ router_logits=hidden_states,
134
+ )
135
+ else:
136
+ router_logits, _ = self.gate(hidden_states)
137
+ final_hidden_states = self.experts(
138
+ hidden_states=hidden_states,
139
+ router_logits=router_logits,
140
+ )
141
+
142
+ if self.is_sequence_parallel and not already_sequence_parallel:
143
+ final_hidden_states = tensor_model_parallel_all_gather(
144
+ final_hidden_states, 0
145
+ )
146
+ final_hidden_states = final_hidden_states[:num_tokens]
147
+
148
+ return final_hidden_states.view(orig_shape)
149
+
150
+
151
+class InternS2MobiusSharedExpertBlock(nn.Module):
152
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
153
+ super().__init__()
154
+
155
+ config = vllm_config.model_config.hf_text_config
156
+ is_sequence_parallel = vllm_config.parallel_config.use_sequence_parallel_moe
157
+ self.shared_expert_gate = ReplicatedLinear(
158
+ config.hidden_size,
159
+ 1,
160
+ bias=False,
161
+ quant_config=None,
162
+ prefix=f"{prefix}.shared_expert_gate",
163
+ )
164
+ self.shared_expert = Qwen2MoeMLP(
165
+ hidden_size=config.hidden_size,
166
+ intermediate_size=config.shared_expert_intermediate_size,
167
+ hidden_act=config.hidden_act,
168
+ quant_config=vllm_config.quant_config,
169
+ expert_gate=self.shared_expert_gate,
170
+ is_sequence_parallel=is_sequence_parallel,
171
+ prefix=f"{prefix}.shared_expert",
172
+ )
173
+
174
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
175
+ return self.shared_expert(hidden_states)
176
+
177
+
178
+class InternS2MobiusDecoderLayer(nn.Module):
179
+ def __init__(
180
+ self,
181
+ vllm_config: VllmConfig,
182
+ layer_type: str,
183
+ prefix: str = "",
184
+ ) -> None:
185
+ super().__init__()
186
+
187
+ config = vllm_config.model_config.hf_text_config
188
+ model_config = vllm_config.model_config
189
+ parallel_config = vllm_config.parallel_config
190
+
191
+ self.layer_type = layer_type
192
+ self.layer_idx = extract_layer_index(prefix)
193
+ self.num_blocks = config.num_blocks
194
+ self.use_attn_reduce_scatter_for_moe = (
195
+ parallel_config.use_sequence_parallel_moe
196
+ and parallel_config.pipeline_parallel_size == 1
197
+ )
198
+
199
+ if self.layer_type == "linear_attention":
200
+ self.linear_attn = QwenGatedDeltaNetAttention(
201
+ config=config,
202
+ vllm_config=vllm_config,
203
+ prefix=f"{prefix}.linear_attn",
204
+ gqa_interleaved_layout=False,
205
+ reduce_results=not self.use_attn_reduce_scatter_for_moe,
206
+ )
207
+ elif self.layer_type == "full_attention":
208
+ self.self_attn = Qwen3NextAttention(
209
+ config,
210
+ model_config=model_config,
211
+ cache_config=vllm_config.cache_config,
212
+ quant_config=vllm_config.quant_config,
213
+ prefix=f"{prefix}.self_attn",
214
+ reduce_results=not self.use_attn_reduce_scatter_for_moe,
215
+ )
216
+ else:
217
+ raise ValueError(f"Invalid layer_type {self.layer_type}")
218
+
219
+ self.mlp = InternS2MobiusSharedExpertBlock(
220
+ vllm_config=vllm_config,
221
+ prefix=f"{prefix}.mlp",
222
+ )
223
+ self.input_layernorm = Qwen3_5RMSNorm(
224
+ config.hidden_size, eps=config.rms_norm_eps
225
+ )
226
+ self.post_attention_layernorm = Qwen3_5RMSNorm(
227
+ config.hidden_size, eps=config.rms_norm_eps
228
+ )
229
+
230
+ self.layer_scale = getattr(config, "layer_scale", False)
231
+ if self.layer_scale:
232
+ self.attn_layer_scale = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
233
+ self.ffn_layer_scale = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
234
+
235
+ def forward(
236
+ self,
237
+ hidden_states: torch.Tensor,
238
+ residual: torch.Tensor | None,
239
+ positions: torch.Tensor,
240
+ meta_mlp: nn.ModuleList,
241
+ ) -> tuple[torch.Tensor, torch.Tensor]:
242
+ full_num_tokens = positions.shape[-1]
243
+ input_is_sequence_parallel = (
244
+ self.use_attn_reduce_scatter_for_moe
245
+ and residual is not None
246
+ and hidden_states.shape[0] != full_num_tokens
247
+ )
248
+
249
+ if residual is None:
250
+ residual = hidden_states
251
+ hidden_states = self.input_layernorm(hidden_states)
252
+ else:
253
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
254
+
255
+ if input_is_sequence_parallel:
256
+ hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
257
+ hidden_states = hidden_states[:full_num_tokens]
258
+
259
+ if self.layer_type == "linear_attention":
260
+ hidden_states = self.linear_attn(hidden_states=hidden_states)
261
+ else:
262
+ hidden_states = self.self_attn(
263
+ hidden_states=hidden_states,
264
+ positions=positions,
265
+ )
266
+
267
+ if self.layer_scale:
268
+ scale = self.attn_layer_scale.to(hidden_states.dtype)
269
+ hidden_states = hidden_states * (scale[0] + 1)
270
+
271
+ if self.use_attn_reduce_scatter_for_moe:
272
+ tp_world_size = get_tensor_model_parallel_world_size()
273
+ sp_pad = (-hidden_states.shape[0]) % tp_world_size
274
+ hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))
275
+ hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)
276
+ if not input_is_sequence_parallel:
277
+ residual = sequence_parallel_chunk(residual)
278
+
279
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
280
+ block_idx = self.layer_idx % self.num_blocks
281
+ routed_output = meta_mlp[block_idx](
282
+ hidden_states,
283
+ already_sequence_parallel=self.use_attn_reduce_scatter_for_moe,
284
+ )
285
+ hidden_states = routed_output + self.mlp(hidden_states)
286
+
287
+ if self.layer_scale:
288
+ scale = self.ffn_layer_scale.to(hidden_states.dtype)
289
+ hidden_states = hidden_states * (scale[0] + 1)
290
+
291
+ return hidden_states, residual
292
+
293
+
294
+@support_torch_compile(
295
+ dynamic_arg_dims={
296
+ "input_ids": 0,
297
+ "positions": -1,
298
+ "intermediate_tensors": 0,
299
+ "inputs_embeds": 0,
300
+ }
301
+)
302
+class InternS2MobiusModel(nn.Module, EagleModelMixin):
303
+ hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper
304
+
305
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
306
+ nn.Module.__init__(self)
307
+
308
+ config = vllm_config.model_config.hf_text_config
309
+ self.config = config
310
+ self.quant_config = vllm_config.quant_config
311
+ self.num_redundant_experts = (
312
+ vllm_config.parallel_config.eplb_config.num_redundant_experts
313
+ )
314
+ self.vocab_size = config.vocab_size
315
+ self.embed_tokens = VocabParallelEmbedding(
316
+ self.vocab_size,
317
+ config.hidden_size,
318
+ )
319
+ self.meta_mlp = nn.ModuleList(
320
+ InternS2MobiusMetaMoeBlock(
321
+ vllm_config=vllm_config,
322
+ prefix=f"{prefix}.meta_mlp.{idx}",
323
+ )
324
+ for idx in range(config.num_blocks)
325
+ )
326
+
327
+ def get_layer(prefix: str) -> InternS2MobiusDecoderLayer:
328
+ return InternS2MobiusDecoderLayer(
329
+ vllm_config,
330
+ layer_type=config.layer_types[extract_layer_index(prefix)],
331
+ prefix=prefix,
332
+ )
333
+
334
+ self.start_layer, self.end_layer, self.layers = make_layers(
335
+ config.num_hidden_layers,
336
+ get_layer,
337
+ prefix=f"{prefix}.layers",
338
+ )
339
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
340
+ ["hidden_states", "residual"], config.hidden_size
341
+ )
342
+
343
+ if get_pp_group().is_last_rank:
344
+ self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
345
+ else:
346
+ self.norm = PPMissingLayer()
347
+
348
+ self.aux_hidden_state_layers: tuple[int, ...] = ()
349
+
350
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
351
+ return self.embed_tokens(input_ids)
352
+
353
+ def forward(
354
+ self,
355
+ input_ids: torch.Tensor | None,
356
+ positions: torch.Tensor,
357
+ intermediate_tensors: IntermediateTensors | None = None,
358
+ inputs_embeds: torch.Tensor | None = None,
359
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
360
+ if get_pp_group().is_first_rank:
361
+ hidden_states = (
362
+ inputs_embeds
363
+ if inputs_embeds is not None
364
+ else self.embed_input_ids(input_ids)
365
+ )
366
+ residual = None
367
+ else:
368
+ assert intermediate_tensors is not None
369
+ hidden_states = intermediate_tensors["hidden_states"]
370
+ residual = intermediate_tensors["residual"]
371
+
372
+ full_num_tokens = positions.shape[-1]
373
+ aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
374
+ for layer_idx, layer in enumerate(
375
+ islice(self.layers, self.start_layer, self.end_layer),
376
+ start=self.start_layer,
377
+ ):
378
+ if (
379
+ hidden_states.shape[0] != full_num_tokens
380
+ and not layer.use_attn_reduce_scatter_for_moe
381
+ ):
382
+ hidden_states, residual = _all_gather_hidden_and_residual(
383
+ hidden_states,
384
+ residual,
385
+ full_num_tokens,
386
+ self.config.hidden_size,
387
+ )
388
+ hidden_states, residual = layer(
389
+ positions=positions,
390
+ hidden_states=hidden_states,
391
+ residual=residual,
392
+ meta_mlp=self.meta_mlp,
393
+ )
394
+ if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[
395
+ 0
396
+ ] != full_num_tokens:
397
+ hidden_states, residual = _all_gather_hidden_and_residual(
398
+ hidden_states,
399
+ residual,
400
+ full_num_tokens,
401
+ self.config.hidden_size,
402
+ )
403
+ self._maybe_add_hidden_state(
404
+ aux_hidden_states,
405
+ layer_idx + 1,
406
+ hidden_states,
407
+ residual,
408
+ )
409
+
410
+ if not get_pp_group().is_last_rank:
411
+ return IntermediateTensors(
412
+ {"hidden_states": hidden_states, "residual": residual}
413
+ )
414
+ if hidden_states.shape[0] != full_num_tokens:
415
+ hidden_states, residual = _all_gather_hidden_and_residual(
416
+ hidden_states,
417
+ residual,
418
+ full_num_tokens,
419
+ self.config.hidden_size,
420
+ )
421
+ hidden_states, _ = self.norm(hidden_states, residual)
422
+ if aux_hidden_states:
423
+ return hidden_states, aux_hidden_states
424
+ return hidden_states
425
+
426
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
427
+ loader = AutoWeightsLoader(self)
428
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
429
+
430
+
431
+class InternS2MobiusForCausalLM(Qwen3_5ForCausalLMBase):
432
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
433
+ config = vllm_config.model_config.hf_text_config
434
+ self.vllm_config = vllm_config
435
+ self.model_config = vllm_config.model_config
436
+ self.quant_config = vllm_config.quant_config
437
+ self.scheduler_config = vllm_config.scheduler_config
438
+
439
+ if vllm_config.cache_config.mamba_cache_mode == "all":
440
+ raise NotImplementedError(
441
+ "Intern-S2-Mobius does not support 'all' prefix caching; "
442
+ "use '--mamba-cache-mode=align' instead."
443
+ )
444
+
445
+ nn.Module.__init__(self)
446
+ self.config = config
447
+ self.model = InternS2MobiusModel(
448
+ vllm_config=vllm_config,
449
+ prefix=maybe_prefix(prefix, "model"),
450
+ )
451
+
452
+ if get_pp_group().is_last_rank:
453
+ if config.tie_word_embeddings:
454
+ self.lm_head = self.model.embed_tokens
455
+ else:
456
+ self.lm_head = ParallelLMHead(
457
+ config.vocab_size,
458
+ config.hidden_size,
459
+ quant_config=self.quant_config,
460
+ prefix=maybe_prefix(prefix, "lm_head"),
461
+ )
462
+ else:
463
+ self.lm_head = PPMissingLayer()
464
+
465
+ self.logits_processor = LogitsProcessor(config.vocab_size)
466
+ self.make_empty_intermediate_tensors = (
467
+ self.model.make_empty_intermediate_tensors
468
+ )
469
+
470
+
471
+class InternS2MobiusMTP(Qwen3_5MoeMTP):
472
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
473
+ speculative_config = vllm_config.speculative_config
474
+ if speculative_config is None or speculative_config.draft_model_config is None:
475
+ raise ValueError(
476
+ "Intern-S2-Mobius MTP requires a draft model configuration."
477
+ )
478
+
479
+ draft_vllm_config = replace(
480
+ vllm_config,
481
+ model_config=speculative_config.draft_model_config,
482
+ )
483
+ with set_current_vllm_config(draft_vllm_config, prefix=prefix):
484
+ super().__init__(vllm_config=draft_vllm_config, prefix=prefix)
485
+
486
+
487
+@MULTIMODAL_REGISTRY.register_processor(
488
+ Qwen3VLMultiModalProcessor,
489
+ info=InternS2MobiusProcessingInfo,
490
+ dummy_inputs=Qwen3VLDummyInputsBuilder,
491
+)
492
+class InternS2MobiusForConditionalGeneration(InternS2PreviewForConditionalGeneration):
493
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
494
+ nn.Module.__init__(self)
495
+
496
+ config = vllm_config.model_config.hf_config
497
+ if getattr(config, "image_token_index", None) is None:
498
+ config.image_token_index = config.image_token_id
499
+ quant_config = vllm_config.quant_config
500
+ multimodal_config = vllm_config.model_config.multimodal_config
501
+
502
+ self.config = config
503
+ self.model_config = vllm_config.model_config
504
+ self.multimodal_config = multimodal_config
505
+ self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
506
+ self.is_multimodal_pruning_enabled = (
507
+ multimodal_config.is_multimodal_pruning_enabled()
508
+ )
509
+ self.video_pruning_rate = multimodal_config.video_pruning_rate
510
+ self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
511
+
512
+ self.use_deepstack = False
513
+
514
+ with self._mark_tower_model(vllm_config, {"image", "video"}):
515
+ self.visual = Qwen3_VisionTransformer(
516
+ config.vision_config,
517
+ norm_eps=getattr(config, "rms_norm_eps", 1e-6),
518
+ quant_config=quant_config,
519
+ prefix=maybe_prefix(prefix, "visual"),
520
+ )
521
+
522
+ with self._mark_language_model(vllm_config):
523
+ self.language_model = InternS2MobiusForCausalLM(
524
+ vllm_config=vllm_config,
525
+ prefix=maybe_prefix(prefix, "language_model"),
526
+ )
527
+
528
+ self.make_empty_intermediate_tensors = (
529
+ self.language_model.make_empty_intermediate_tensors
530
+ )
531
+ self.set_moe_parameters()
532
+
533
+ def set_moe_parameters(self) -> None:
534
+ meta_mlp = self.language_model.model.meta_mlp
535
+ if not meta_mlp:
536
+ raise RuntimeError("Intern-S2-Mobius requires at least one MoE bank.")
537
+
538
+ self.moe_layers = [moe.experts for moe in meta_mlp]
539
+ example_moe = meta_mlp[0]
540
+ self.num_moe_layers = len(self.moe_layers)
541
+ self.num_expert_groups = 1
542
+ self.num_shared_experts = 0
543
+ self.num_logical_experts = example_moe.n_logical_experts
544
+ self.num_physical_experts = example_moe.n_physical_experts
545
+ self.num_local_physical_experts = example_moe.n_local_physical_experts
546
+ self.num_routed_experts = example_moe.n_routed_experts
547
+ self.num_redundant_experts = example_moe.n_redundant_experts
548
+
549
+ def update_physical_experts_metadata(
550
+ self,
551
+ num_physical_experts: int,
552
+ num_local_physical_experts: int,
553
+ ) -> None:
554
+ assert self.num_local_physical_experts == num_local_physical_experts
555
+ self.num_physical_experts = num_physical_experts
556
+ self.num_local_physical_experts = num_local_physical_experts
557
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
558
+ for moe in self.language_model.model.meta_mlp:
559
+ moe.n_local_physical_experts = num_local_physical_experts
560
+ moe.n_physical_experts = num_physical_experts
561
+ moe.n_redundant_experts = self.num_redundant_experts
562
+ moe.experts.update_expert_map()
vllm/model_executor/models/registry.py
+5
@@ -447,6 +447,10 @@ _MULTIMODAL_MODELS = {
447
"interns1_pro",
448
"InternS1ProForConditionalGeneration",
449
),
450
+ "InternS2MobiusForConditionalGeneration": (
451
+ "interns2_mobius",
452
+ "InternS2MobiusForConditionalGeneration",
453
+ ),
454
"InternS2PreviewForConditionalGeneration": (
455
"interns2_preview",
456
"InternS2PreviewForConditionalGeneration",
@@ -661,6 +665,7 @@ _SPECULATIVE_DECODING_MODELS = {
665
"Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),
666
"Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
667
"Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
668
+ "InternS2MobiusMTP": ("interns2_mobius", "InternS2MobiusMTP"),
669
"HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"),
670
"KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"),
671
# Temporarily disabled.