From 94a7c5cc02687e53f402911458a85a65d54b3f8b Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 6 Jan 2026 16:44:10 +0800 Subject: [PATCH 01/24] support interns1 pro Signed-off-by: zxy --- examples/offline_inference/vision_language.py | 35 + .../layers/rotary_embedding/__init__.py | 22 + .../layers/rotary_embedding/fope.py | 149 ++++ vllm/model_executor/models/interns1_1.py | 806 ++++++++++++++++++ vllm/model_executor/models/qwen3_vl.py | 6 +- vllm/model_executor/models/registry.py | 4 + 6 files changed, 1021 insertions(+), 1 deletion(-) create mode 100644 vllm/model_executor/layers/rotary_embedding/fope.py create mode 100644 vllm/model_executor/models/interns1_1.py diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index dfca7d5c9c9a..a766a763edf6 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -729,6 +729,40 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData: ) +# Intern-S1-1 +def run_interns1_1(questions: list[str], modality: str) -> ModelRequestData: + model_name = "internlm-private/InternS1_1_1T_A22_1217" + + engine_args = EngineArgs( + model=model_name, + trust_remote_code=True, + max_model_len=8192, + max_num_seqs=2, + limit_mm_per_prompt={modality: 1}, + enforce_eager=True, + tensor_parallel_size=4, + ) + + if modality == "image": + placeholder = "<|image_pad|>" + elif modality == "video": + placeholder = "<|video_pad|>" + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + messages = [ + [{"role": "user", "content": f"{placeholder}\n{question}"}] + for question in questions + ] + prompts = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # InternVL def run_internvl(questions: list[str], modality: str) -> ModelRequestData: model_name = "OpenGVLab/InternVL3-2B" @@ -1844,6 +1878,7 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData: "hyperclovax_seed_vision": run_hyperclovax_seed_vision, "idefics3": run_idefics3, "interns1": run_interns1, + "interns1_1": run_interns1_1, "internvl_chat": run_internvl, "keye_vl": run_keye_vl, "keye_vl1_5": run_keye_vl1_5, diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index 452b87ea4e7a..5db30834b480 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -11,6 +11,7 @@ from .dual_chunk_rope import DualChunkRotaryEmbedding from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding +from .fope import FourierRotaryEmbedding from .linear_scaling_rope import LinearScalingRotaryEmbedding from .llama3_rope import Llama3RotaryEmbedding from .llama4_vision_rope import Llama4VisionRotaryEmbedding @@ -101,6 +102,27 @@ def get_rope( mrope_section=rope_parameters["mrope_section"], mrope_interleaved=rope_parameters.get("mrope_interleaved", False), ) + elif "use_fope" in rope_parameters and rope_parameters["use_fope"]: + extra_kwargs = { + k: v + for k, v in rope_parameters.items() + if k + in ( + "num_key_value_heads", + "num_inv_freq", + "fope_sep_head", + "fope_init_factor", + ) + } + rotary_emb = FourierRotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + **extra_kwargs, + ) else: rotary_emb = RotaryEmbedding( head_size, diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py new file mode 100644 index 000000000000..127c0a464f4c --- /dev/null +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +import torch.nn.functional as F +from torch import nn + +from .base import RotaryEmbedding + + +class FourierRotaryEmbedding(RotaryEmbedding): + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + is_neox_style: bool, + dtype: torch.dtype, + # extra parameters for FoPE + num_key_value_heads: int, + num_inv_freq: int, + fope_sep_head: bool, + fope_init_factor: float, + ): + # fope related parameters + self.num_key_value_heads = num_key_value_heads + self.num_inv_freq = num_inv_freq + self.fope_sep_head = fope_sep_head + self.fope_init_factor = fope_init_factor + + # init inv_freq + inv_freq = 1.0 / ( + base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim) + ) + inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) + if self.num_inv_freq is not None: + num_inv_freq = self.num_inv_freq + inv_freq_idx_selected[num_inv_freq:] = False + else: + inv_freq_idx_selected = inv_freq > ( + 2.0 * torch.pi / max_position_embeddings + ) + num_inv_freq = inv_freq_idx_selected.sum().item() + inv_freq = inv_freq[inv_freq_idx_selected] + self.inv_freq = inv_freq + + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + # setup parameters + self.input_dim = self.inv_freq.shape[-1] + self.output_dim = self.inv_freq.shape[-1] + self.cos_coef = nn.Parameter( + torch.empty(num_key_value_heads, self.input_dim, self.output_dim), + requires_grad=False, + ) + self.sin_coef = nn.Parameter( + torch.empty(num_key_value_heads, self.input_dim, self.output_dim), + requires_grad=False, + ) + + def _compute_inv_freq(self, base: float) -> torch.Tensor: + return None + + def _compute_cos_sin_cache(self) -> torch.Tensor: + # FIXME: zhouxinyu, implement FoPE cos/sin cache computation + return torch.zeros(1) + + def get_step_eye(self, _param): + import math + + _step_eye = torch.zeros_like(_param) + + step = math.ceil(self.input_dim / self.output_dim) + for i in range(self.output_dim): + if i * step < self.input_dim: + _step_eye[..., i * step, i] = 1.0 + + return _step_eye + + def forward_native(self, x: torch.Tensor, positions: torch.Tensor): + # expand x, positions to additional batch size dim + if x.dim() == 2: + # (seq_len, hidden_size) -> (bsz=1, seq_len, hidden_size) + x = x.unsqueeze(0) + if positions.dim() == 1: + # (seq_len) -> (bsz=1, seq_len) + positions = positions.unsqueeze(0) + + # Core RoPE block + + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(positions.shape[0], -1, 1) + ) # (40) -> (1, 40, 1) -> (bsz, 40, 1) + position_ids_expanded = positions[ + :, None, : + ].float() # (bsz, seq_len) -> (bsz, 1, seq_len) + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + batch_size, seq_len, hidden_size = x.shape + with torch.autocast(device_type=device_type, enabled=False): + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(1, 2) + if self.fope_sep_head: + pos_cos = ( + freqs.cos() + .unsqueeze(1) + .expand(batch_size, self.num_key_value_heads, seq_len, -1) + ) + pos_sin = ( + freqs.sin() + .unsqueeze(1) + .expand(batch_size, self.num_key_value_heads, seq_len, -1) + ) + else: + pos_cos = freqs.cos() + pos_sin = freqs.sin() + + if self.fope_sep_head: + # (1, 1, 8192, 40) x (1, 40, 40) -> (1, 8192, 1, 40) + sin = torch.einsum("bhtD, hDd -> bthd", pos_sin, self.sin_coef.float()) + cos = torch.einsum("bhtD, hDd -> bthd", pos_cos, self.cos_coef.float()) + else: + sin = torch.einsum("btD, Dd -> btd", pos_sin, self.sin_coef.float()) + cos = torch.einsum("btD, Dd -> btd", 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) + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + def forward_cuda(self, x: torch.Tensor, positions: torch.Tensor): + # FIXME, zhouxinyu, implement FoPE cuda forward computation + return self.forward_native(x, positions) diff --git a/vllm/model_executor/models/interns1_1.py b/vllm/model_executor/models/interns1_1.py new file mode 100644 index 000000000000..fc0c511e2399 --- /dev/null +++ b/vllm/model_executor/models/interns1_1.py @@ -0,0 +1,806 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2025 The vLLM team. +# Copyright 2025 The Qwen Team. +# Copyright 2025 The HuggingFace Inc. team. +# All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only InternS1_1 model compatible with HuggingFace weights.""" + +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import torch +from torch import nn +from transformers import AutoProcessor, PretrainedConfig + +from vllm.attention.layer import Attention +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.common import ( + ApplyRotaryEmb, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from .interfaces import MixtureOfExperts +from .qwen3_moe import ( + Qwen3MoeForCausalLM, + Qwen3MoeMLP, + Qwen3MoeSparseMoeBlock, +) +from .qwen3_vl import ( + Qwen3_VisionTransformer, + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) +from .utils import ( + extract_layer_index, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) + +logger = init_logger(__name__) + + +class InternS1_1_ProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config() + + def get_hf_processor(self, **kwargs: object) -> AutoProcessor: + return AutoProcessor.from_pretrained( + self.ctx.model_config.model, + trust_remote_code=True, + **kwargs, + ) + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def apply_rotary_pos_emb_sep(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + num_groups = int(q.shape[unsqueeze_dim] // cos.shape[unsqueeze_dim]) + cos_rep = repeat_kv(cos, num_groups) + sin_rep = repeat_kv(sin, num_groups) + q_embed = (q * cos_rep) + (rotate_half(q) * sin_rep) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Qwen3MoeAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int | None = None, + rms_norm_eps: float = 1e-06, + qkv_bias: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim or (hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + # qkv: (8192, 1280) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # q: (seq_len, q_size), (8192, 1024) + # k: (seq_len, kv_size), (8192, 128) + # v: (seq_len, kv_size), (8192, 128) + + # Add qk-norm + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) + q_by_head = self.q_norm(q_by_head) + q = q_by_head.view(q.shape) + + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) + k_by_head = self.k_norm(k_by_head) + k = k_by_head.view(k.shape) + # q, k = self.rotary_emb(positions, q, k) + + # hidden_states: (seq_len, hidden_size), (8192, 2048) + # q: (seq_len, q_size), (8192, 1024) + # k: (seq_len, kv_size), (8192, 128) + # v: (seq_len, kv_size), (8192, 128) + # cos: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) + # sin: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) + + # apply rotary embedding sep head + q = q.unflatten( + -1, (self.num_heads, self.head_dim) + ) # (seq_len, num_heads, head_dim), (8192, 8, 128) + k = k.unflatten( + -1, (self.num_kv_heads, self.head_dim) + ) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) + v = v.unflatten( + -1, (self.num_kv_heads, self.head_dim) + ) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) + q = q.unsqueeze(0) # (bsz, seq_len, num_heads, head_dim), (1, 8192, 8, 128) + k = k.unsqueeze(0) # (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) + v = v.unsqueeze(0) # (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) + # cos, sin: (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) + cos = rotary_pos_emb_cos + sin = rotary_pos_emb_sin + q, k = apply_rotary_pos_emb_sep( + q, + k, + cos, + sin, + ) + q = q.squeeze(0) # (seq_len, num_heads, head_dim), (8192, 8, 128) + k = k.squeeze(0) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) + v = v.squeeze(0) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) + q = q.flatten(1, 2) # (seq_len, q_size), (8192, 1024) + k = k.flatten(1, 2) # (seq_len, kv_size), (8192, 128) + v = v.flatten(1, 2) # (seq_len, kv_size), (8192, 128) + + attn_output = self.attn(q, k, v) # (seq_len, q_size), (8192, 1024) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen3MoeDecoderLayer(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config = vllm_config.model_config.hf_text_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.hidden_size = config.hidden_size + self.self_attn = Qwen3MoeAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=getattr(config, "attention_bias", False), + head_dim=getattr(config, "head_dim", None), + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + + # `mlp_only_layers` in the config. + layer_idx = extract_layer_index(prefix) + mlp_only_layers = ( + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers + ) + if (layer_idx not in mlp_only_layers) and ( + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 + ): + self.mlp = Qwen3MoeSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.mlp" + ) + else: + self.mlp = Qwen3MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + hidden_states=hidden_states, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl, + # otherwise (seq_len, ). + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class Qwen3MoeLLMModel(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + eplb_config = parallel_config.eplb_config + self.num_redundant_experts = eplb_config.num_redundant_experts + + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.config = config + self.quant_config = quant_config + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + + # build rotary embedding + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + tp_size = get_tensor_model_parallel_world_size() + num_key_value_heads = max(1, config.num_key_value_heads // tp_size) + dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = config.hidden_size // config.num_attention_heads + + rope_scaling = config.rope_scaling + fope_init_factor = rope_scaling.get("fope_init_factor", None) + fope_sep_head = rope_scaling.get("fope_sep_head", None) + num_inv_freq = rope_scaling.get("num_inv_freq", None) + + 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) + rope_parameters = { + "rope_theta": config.rope_theta, + "num_key_value_heads": num_key_value_heads, + "use_fope": use_fope, + "fope_init_factor": fope_init_factor, + "fope_sep_head": fope_sep_head, + "num_inv_freq": num_inv_freq, + } + self.rotary_emb = get_rope( + head_size=head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + + # build layers + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: Qwen3MoeDecoderLayer(vllm_config=vllm_config, prefix=prefix), + prefix=f"{prefix}.layers", + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + # Track layers for auxiliary hidden state outputs (EAGLE3) + self.aux_hidden_state_layers: tuple[int, ...] = () + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + deepstack_input_embeds: IntermediateTensors | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + # rotary embedding + rotary_pos_emb_cos, rotary_pos_emb_sin = self.rotary_emb( + hidden_states, positions + ) + + # decoder layers + for layer_idx, layer in islice( + enumerate(self.layers), self.start_layer, self.end_layer + ): + hidden_states, residual = layer( + hidden_states=hidden_states, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + residual=residual, + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + + def load_fused_expert_weights( + self, + name: str, + params_dict: dict, + loaded_weight: torch.Tensor, + shard_id: str, + num_experts: int, + ) -> bool: + param = params_dict[name] + weight_loader = typing.cast(Callable[..., bool], param.weight_loader) + loaded_local_expert = False + for expert_id in range(num_experts): + curr_expert_weight = loaded_weight[expert_id] + success = weight_loader( + param, + curr_expert_weight, + name, + shard_id, + expert_id, + return_success=True, + ) + if success: + loaded_local_expert = True + + return loaded_local_expert + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + # Skip loading extra parameters for GPTQ/modelopt models. + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + expert_params_mapping = self.get_expert_mapping() + is_fused_expert = False + fused_expert_params_mapping = [ + ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), + ("experts.w2_weight", "experts.down_proj", 0, "w2"), + ] + num_experts = self.config.num_experts + for name, loaded_weight in weights: + for param_name, weight_name, shard_id in stacked_params_mapping: + if "experts.gate_up_proj" in name or "experts.down_proj" in name: + is_fused_expert = True + expert_params_mapping = fused_expert_params_mapping + + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra parameters for GPTQ/modelopt models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + # Skip layers on other devices. + if is_pp_missing_parameter(name, self): + continue + if name.endswith("scale"): + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + # Anyway, this is an expert weight and should not be + # attempted to load as other weights later + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + if is_fused_expert: + loaded_weight = loaded_weight.transpose(-1, -2) # no bias + if "experts.gate_up_proj" in name: + loaded_weight = loaded_weight.chunk(2, dim=-2) + success_w1 = self.load_fused_expert_weights( + name_mapped, + params_dict, + loaded_weight[0], + "w1", + num_experts, + ) + success_w3 = self.load_fused_expert_weights( + name_mapped, + params_dict, + loaded_weight[1], + "w3", + num_experts, + ) + success = success_w1 and success_w3 + else: + # down_proj + success = self.load_fused_expert_weights( + name_mapped, + params_dict, + loaded_weight, + shard_id, + num_experts, + ) + else: + # Skip loading extra parameters for GPTQ/modelopt models + if ( + name_mapped.endswith(ignore_suffixes) + and name_mapped not in params_dict + ): + continue + param = params_dict[name_mapped] + # We should ask the weight loader to return success or + # not here since otherwise we may skip experts with + # other available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + else: + if is_expert_weight: + # We've checked that this is an expert weight + # However it's not mapped locally to this rank + # So we simply skip it + continue + # Skip loading extra parameters for GPTQ/modelopt models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + # Skip layers on other devices. + if is_pp_missing_parameter(name, self): + continue + # Remapping the name of FP8 kv-scale. + if name.endswith("kv_scale"): + remapped_kv_scale_name = name.replace( + ".kv_scale", ".attn.kv_scale" + ) + if remapped_kv_scale_name not in params_dict: + logger.warning_once( + "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.", # noqa: E501 + name, + remapped_kv_scale_name, + ) + continue + else: + name = remapped_kv_scale_name + param = params_dict[name] + # FIXME: zhouxinyu, put this logic as FoPE weight_loader + if name == "rotary_emb.cos_coef" or name == "rotary_emb.sin_coef": + from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + + world_size = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + num_key_value_heads = loaded_weight.size(0) + + if num_key_value_heads < world_size: + n_replicate = world_size // num_key_value_heads + world_size = num_key_value_heads + rank = rank // n_replicate + + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] + param.copy_(loaded_weight) + else: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super(Qwen3MoeForCausalLM, self).__init__() + self.config = vllm_config.model_config.hf_config.text_config + self.quant_config = vllm_config.quant_config + self.model = Qwen3MoeLLMModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_processor = LogitsProcessor(self.config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + +class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts): + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.language_model.model.layers: + if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + def set_moe_parameters(self): + self.expert_weights = [] + + self.moe_layers = [] + example_moe = None + for layer in self.language_model.model.layers: + if hasattr(layer, "mlp") and isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + example_moe = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_moe is None: + raise RuntimeError("No Qwen3Moe layer found in the language_model.") + + # Set MoE hyperparameters + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=InternS1_1_ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class InternS1_1_ForConditionalGeneration( + Qwen3VLForConditionalGeneration, Qwen3VLMoeMixtureOfExperts +): + is_3d_moe_weight: bool = True + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super(Qwen3VLForConditionalGeneration, self).__init__() + config: PretrainedConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.config = config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self.video_pruning_rate = multimodal_config.video_pruning_rate + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + + if not multimodal_config.get_limit_per_prompt( + "image" + ) and not multimodal_config.get_limit_per_prompt("video"): + self.visual = None + else: + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + multimodal_config=multimodal_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + self.language_model = Qwen3MoeLLMForCausalLM( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") + ) + # Whether to include the gate_up_proj mapping is determined by + # the language model. + self.packed_modules_mapping = ( + self.packed_modules_mapping | self.language_model.packed_modules_mapping + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + # Set MoE hyperparameters + self.set_moe_parameters() diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index f47062c101e3..e5866a4ff95e 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -332,7 +332,11 @@ def __init__( self.spatial_merge_size = vision_config.spatial_merge_size self.spatial_merge_unit = self.spatial_merge_size**2 self.temporal_patch_size = vision_config.temporal_patch_size - self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes + self.deepstack_visual_indexes = ( + vision_config.deepstack_visual_indexes + if hasattr(vision_config, "deepstack_visual_indexes") + else [] + ) self.num_grid_per_side = int(self.num_position_embeddings**0.5) # NOTE: This is used for creating empty tensor for all_gather for diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 6cbca9a69b3d..9e81e6848f33 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -329,6 +329,10 @@ "interns1", "InternS1ForConditionalGeneration", ), + "InternS1_1_ForConditionalGeneration": ( + "interns1_1", + "InternS1_1_ForConditionalGeneration", + ), "Idefics3ForConditionalGeneration": ( "idefics3", "Idefics3ForConditionalGeneration", From 6aa82db32777a99afbe680233379b892be985958 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 6 Jan 2026 18:48:24 +0800 Subject: [PATCH 02/24] better weight loading Signed-off-by: zxy --- .../layers/rotary_embedding/fope.py | 32 +++++++++++++++++++ vllm/model_executor/models/interns1_1.py | 27 +++------------- 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index 127c0a464f4c..e70be8e00590 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -3,10 +3,30 @@ import torch import torch.nn.functional as F from torch import nn +from torch.nn.parameter import Parameter from .base import RotaryEmbedding +def fope_coef_weight_loader(param: Parameter, loaded_weight: torch.Tensor): + from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + + world_size = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + num_key_value_heads = loaded_weight.size(0) + + if num_key_value_heads < world_size: + n_replicate = world_size // num_key_value_heads + world_size = num_key_value_heads + rank = rank // n_replicate + + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] + param.copy_(loaded_weight) + + class FourierRotaryEmbedding(RotaryEmbedding): def __init__( self, @@ -59,14 +79,26 @@ def __init__( torch.empty(num_key_value_heads, self.input_dim, self.output_dim), requires_grad=False, ) + self.cos_coef.weight_loader = ( + lambda param, loaded_weight: fope_coef_weight_loader(param, loaded_weight) + ) + self.sin_coef.weight_loader = ( + lambda param, loaded_weight: fope_coef_weight_loader(param, loaded_weight) + ) def _compute_inv_freq(self, base: float) -> torch.Tensor: + """Compute the inverse frequency.""" return None def _compute_cos_sin_cache(self) -> torch.Tensor: + """Compute the cos and sin cache.""" # FIXME: zhouxinyu, implement FoPE cos/sin cache computation return torch.zeros(1) + def apply_rotary_emb(self): + """Customized apply_rotary_emb function for FoPE.""" + pass + def get_step_eye(self, _param): import math diff --git a/vllm/model_executor/models/interns1_1.py b/vllm/model_executor/models/interns1_1.py index fc0c511e2399..7d8e32e6052f 100644 --- a/vllm/model_executor/models/interns1_1.py +++ b/vllm/model_executor/models/interns1_1.py @@ -645,29 +645,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: else: name = remapped_kv_scale_name param = params_dict[name] - # FIXME: zhouxinyu, put this logic as FoPE weight_loader - if name == "rotary_emb.cos_coef" or name == "rotary_emb.sin_coef": - from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - ) - - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - num_key_value_heads = loaded_weight.size(0) - - if num_key_value_heads < world_size: - n_replicate = world_size // num_key_value_heads - world_size = num_key_value_heads - rank = rank // n_replicate - - loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] - param.copy_(loaded_weight) - else: - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params From 668455ac14c5815728e57cddaefa92b28f7d5b46 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 6 Jan 2026 20:43:12 +0800 Subject: [PATCH 03/24] better inv freq calc location Signed-off-by: zxy --- .../layers/rotary_embedding/fope.py | 38 ++++++++++--------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index e70be8e00590..66592cee43ce 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -48,22 +48,6 @@ def __init__( self.fope_sep_head = fope_sep_head self.fope_init_factor = fope_init_factor - # init inv_freq - inv_freq = 1.0 / ( - base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim) - ) - inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) - if self.num_inv_freq is not None: - num_inv_freq = self.num_inv_freq - inv_freq_idx_selected[num_inv_freq:] = False - else: - inv_freq_idx_selected = inv_freq > ( - 2.0 * torch.pi / max_position_embeddings - ) - num_inv_freq = inv_freq_idx_selected.sum().item() - inv_freq = inv_freq[inv_freq_idx_selected] - self.inv_freq = inv_freq - super().__init__( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype ) @@ -88,10 +72,30 @@ def __init__( def _compute_inv_freq(self, base: float) -> torch.Tensor: """Compute the inverse frequency.""" - return None + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim + ) + ) + + inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) + if self.num_inv_freq is not None: + num_inv_freq = self.num_inv_freq + inv_freq_idx_selected[num_inv_freq:] = False + else: + inv_freq_idx_selected = inv_freq > ( + 2.0 * torch.pi / self.max_position_embeddings + ) + num_inv_freq = inv_freq_idx_selected.sum().item() + + 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.""" + self.inv_freq = self._compute_inv_freq(self.base) # FIXME: zhouxinyu, implement FoPE cos/sin cache computation return torch.zeros(1) From 69171b126637547b784c0501e4ce9f750b066240 Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 7 Jan 2026 18:06:55 +0800 Subject: [PATCH 04/24] adapt to new name Signed-off-by: zxy --- examples/offline_inference/vision_language.py | 6 +- .../models/{interns1_1.py => interns1_pro.py} | 72 ++++--------------- vllm/model_executor/models/registry.py | 6 +- 3 files changed, 21 insertions(+), 63 deletions(-) rename vllm/model_executor/models/{interns1_1.py => interns1_pro.py} (91%) diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index a766a763edf6..386eb6c2fd19 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -729,8 +729,8 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData: ) -# Intern-S1-1 -def run_interns1_1(questions: list[str], modality: str) -> ModelRequestData: +# Intern-S1-Pro +def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData: model_name = "internlm-private/InternS1_1_1T_A22_1217" engine_args = EngineArgs( @@ -1878,7 +1878,7 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData: "hyperclovax_seed_vision": run_hyperclovax_seed_vision, "idefics3": run_idefics3, "interns1": run_interns1, - "interns1_1": run_interns1_1, + "interns1_pro": run_interns1_pro, "internvl_chat": run_internvl, "keye_vl": run_keye_vl, "keye_vl1_5": run_keye_vl1_5, diff --git a/vllm/model_executor/models/interns1_1.py b/vllm/model_executor/models/interns1_pro.py similarity index 91% rename from vllm/model_executor/models/interns1_1.py rename to vllm/model_executor/models/interns1_pro.py index 7d8e32e6052f..67789f2b8afd 100644 --- a/vllm/model_executor/models/interns1_1.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -63,7 +63,6 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors -from .interfaces import MixtureOfExperts from .qwen3_moe import ( Qwen3MoeForCausalLM, Qwen3MoeMLP, @@ -76,6 +75,7 @@ Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, ) +from .qwen3_vl_moe import Qwen3VLMoeMixtureOfExperts from .utils import ( extract_layer_index, is_pp_missing_parameter, @@ -87,7 +87,7 @@ logger = init_logger(__name__) -class InternS1_1_ProcessingInfo(Qwen3VLProcessingInfo): +class InternS1ProProcessingInfo(Qwen3VLProcessingInfo): def get_hf_config(self): return self.ctx.get_hf_config() @@ -115,7 +115,7 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) -def apply_rotary_pos_emb_sep(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): +def apply_rotary_pos_emb_sep(q, k, cos, sin, unsqueeze_dim=1): num_groups = int(q.shape[unsqueeze_dim] // cos.shape[unsqueeze_dim]) cos_rep = repeat_kv(cos, num_groups) sin_rep = repeat_kv(sin, num_groups) @@ -124,7 +124,7 @@ def apply_rotary_pos_emb_sep(q, k, cos, sin, position_ids=None, unsqueeze_dim=1) return q_embed, k_embed -class Qwen3MoeAttention(nn.Module): +class InternS1ProMoeAttention(nn.Module): def __init__( self, hidden_size: int, @@ -255,7 +255,7 @@ def forward( return output -class Qwen3MoeDecoderLayer(nn.Module): +class InternS1ProMoeDecoderLayer(nn.Module): def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -264,7 +264,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: quant_config = vllm_config.quant_config self.hidden_size = config.hidden_size - self.self_attn = Qwen3MoeAttention( + self.self_attn = InternS1ProMoeAttention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, @@ -328,14 +328,12 @@ def forward( @support_torch_compile( dynamic_arg_dims={ "input_ids": 0, - # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl, - # otherwise (seq_len, ). "positions": -1, "intermediate_tensors": 0, "inputs_embeds": 0, } ) -class Qwen3MoeLLMModel(nn.Module): +class InternS1ProMoeLLMModel(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -392,7 +390,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # build layers self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, - lambda prefix: Qwen3MoeDecoderLayer(vllm_config=vllm_config, prefix=prefix), + lambda prefix: InternS1ProMoeDecoderLayer( + vllm_config=vllm_config, prefix=prefix + ), prefix=f"{prefix}.layers", ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -653,12 +653,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params -class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM): +class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super(Qwen3MoeForCausalLM, self).__init__() self.config = vllm_config.model_config.hf_config.text_config self.quant_config = vllm_config.quant_config - self.model = Qwen3MoeLLMModel( + self.model = InternS1ProMoeLLMModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.lm_head = ParallelLMHead( @@ -675,54 +675,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) -class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts): - def update_physical_experts_metadata( - self, - num_physical_experts: int, - num_local_physical_experts: int, - ) -> None: - assert self.num_local_physical_experts == num_local_physical_experts - self.num_physical_experts = num_physical_experts - self.num_local_physical_experts = num_local_physical_experts - self.num_redundant_experts = num_physical_experts - self.num_logical_experts - for layer in self.language_model.model.layers: - if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): - moe = layer.mlp - moe.n_local_physical_experts = num_local_physical_experts - moe.n_physical_experts = num_physical_experts - moe.n_redundant_experts = self.num_redundant_experts - moe.experts.update_expert_map() - - def set_moe_parameters(self): - self.expert_weights = [] - - self.moe_layers = [] - example_moe = None - for layer in self.language_model.model.layers: - if hasattr(layer, "mlp") and isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): - example_moe = layer.mlp - self.moe_layers.append(layer.mlp.experts) - - if example_moe is None: - raise RuntimeError("No Qwen3Moe layer found in the language_model.") - - # Set MoE hyperparameters - self.num_moe_layers = len(self.moe_layers) - self.num_expert_groups = 1 - self.num_shared_experts = 0 - self.num_logical_experts = example_moe.n_logical_experts - self.num_physical_experts = example_moe.n_physical_experts - self.num_local_physical_experts = example_moe.n_local_physical_experts - self.num_routed_experts = example_moe.n_routed_experts - self.num_redundant_experts = example_moe.n_redundant_experts - - @MULTIMODAL_REGISTRY.register_processor( Qwen3VLMultiModalProcessor, - info=InternS1_1_ProcessingInfo, + info=InternS1ProProcessingInfo, dummy_inputs=Qwen3VLDummyInputsBuilder, ) -class InternS1_1_ForConditionalGeneration( +class InternS1ProForConditionalGeneration( Qwen3VLForConditionalGeneration, Qwen3VLMoeMixtureOfExperts ): is_3d_moe_weight: bool = True @@ -761,7 +719,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "visual"), ) - self.language_model = Qwen3MoeLLMForCausalLM( + self.language_model = InternS1ProMoeLLMForCausalLM( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") ) # Whether to include the gate_up_proj mapping is determined by diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5096c3948a84..cb2c1c36de91 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -331,9 +331,9 @@ "interns1", "InternS1ForConditionalGeneration", ), - "InternS1_1_ForConditionalGeneration": ( - "interns1_1", - "InternS1_1_ForConditionalGeneration", + "InternS1ProForConditionalGeneration": ( + "interns1_pro", + "InternS1ProForConditionalGeneration", ), "Idefics3ForConditionalGeneration": ( "idefics3", From ced69649969baf916cb659568935f1ad7c4d6c74 Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 7 Jan 2026 18:11:28 +0800 Subject: [PATCH 05/24] minor Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 67789f2b8afd..7fc2bc184cc8 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -22,7 +22,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Inference-only InternS1_1 model compatible with HuggingFace weights.""" +"""Inference-only InternS1Pro model compatible with HuggingFace weights.""" import typing from collections.abc import Callable, Iterable From 44ef618b89876d689e372746f7f76c05df34c5f1 Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 7 Jan 2026 21:17:50 +0800 Subject: [PATCH 06/24] tiny cleanups Signed-off-by: zxy --- .../layers/rotary_embedding/fope.py | 54 +++++++++---------- vllm/model_executor/models/interns1_pro.py | 53 +++++++----------- 2 files changed, 43 insertions(+), 64 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index 66592cee43ce..071e374be56b 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -8,25 +8,6 @@ from .base import RotaryEmbedding -def fope_coef_weight_loader(param: Parameter, loaded_weight: torch.Tensor): - from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - ) - - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - num_key_value_heads = loaded_weight.size(0) - - if num_key_value_heads < world_size: - n_replicate = world_size // num_key_value_heads - world_size = num_key_value_heads - rank = rank // n_replicate - - loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] - param.copy_(loaded_weight) - - class FourierRotaryEmbedding(RotaryEmbedding): def __init__( self, @@ -63,12 +44,8 @@ def __init__( torch.empty(num_key_value_heads, self.input_dim, self.output_dim), requires_grad=False, ) - self.cos_coef.weight_loader = ( - lambda param, loaded_weight: fope_coef_weight_loader(param, loaded_weight) - ) - self.sin_coef.weight_loader = ( - lambda param, loaded_weight: fope_coef_weight_loader(param, loaded_weight) - ) + self.cos_coef.weight_loader = self.fope_coef_weight_loader + self.sin_coef.weight_loader = self.fope_coef_weight_loader def _compute_inv_freq(self, base: float) -> torch.Tensor: """Compute the inverse frequency.""" @@ -96,7 +73,7 @@ def _compute_inv_freq(self, base: float) -> torch.Tensor: def _compute_cos_sin_cache(self) -> torch.Tensor: """Compute the cos and sin cache.""" self.inv_freq = self._compute_inv_freq(self.base) - # FIXME: zhouxinyu, implement FoPE cos/sin cache computation + # TODO: zhouxinyu, implement FoPE cos/sin cache computation return torch.zeros(1) def apply_rotary_emb(self): @@ -118,14 +95,13 @@ def get_step_eye(self, _param): def forward_native(self, x: torch.Tensor, positions: torch.Tensor): # expand x, positions to additional batch size dim if x.dim() == 2: - # (seq_len, hidden_size) -> (bsz=1, seq_len, hidden_size) + # (seq_len, hidden_size) -> (bsz, seq_len, hidden_size) x = x.unsqueeze(0) if positions.dim() == 1: - # (seq_len) -> (bsz=1, seq_len) + # (seq_len) -> (bsz, seq_len) positions = positions.unsqueeze(0) # Core RoPE block - inv_freq_expanded = ( self.inv_freq[None, :, None].float().expand(positions.shape[0], -1, 1) ) # (40) -> (1, 40, 1) -> (bsz, 40, 1) @@ -181,5 +157,23 @@ def forward_native(self, x: torch.Tensor, positions: torch.Tensor): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def forward_cuda(self, x: torch.Tensor, positions: torch.Tensor): - # FIXME, zhouxinyu, implement FoPE cuda forward computation + # TODO, zhouxinyu, implement FoPE cuda forward computation return self.forward_native(x, positions) + + def fope_coef_weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + + world_size = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + num_key_value_heads = loaded_weight.size(0) + + if num_key_value_heads < world_size: + n_replicate = world_size // num_key_value_heads + world_size = num_key_value_heads + rank = rank // n_replicate + + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] + param.copy_(loaded_weight) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 7fc2bc184cc8..249a45738580 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -49,9 +49,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.rotary_embedding.common import ( - ApplyRotaryEmb, -) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -189,7 +186,8 @@ def __init__( self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) + # TODO: zhouxinyu, use vllm ApplyRotaryEmb + # self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) def forward( self, @@ -212,45 +210,32 @@ def forward( k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) k_by_head = self.k_norm(k_by_head) k = k_by_head.view(k.shape) - # q, k = self.rotary_emb(positions, q, k) # hidden_states: (seq_len, hidden_size), (8192, 2048) # q: (seq_len, q_size), (8192, 1024) - # k: (seq_len, kv_size), (8192, 128) - # v: (seq_len, kv_size), (8192, 128) - # cos: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) - # sin: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) + # k, v: (seq_len, kv_size), (8192, 128) + # cos, sin: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) # apply rotary embedding sep head - q = q.unflatten( - -1, (self.num_heads, self.head_dim) - ) # (seq_len, num_heads, head_dim), (8192, 8, 128) - k = k.unflatten( - -1, (self.num_kv_heads, self.head_dim) - ) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) - v = v.unflatten( - -1, (self.num_kv_heads, self.head_dim) - ) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) - q = q.unsqueeze(0) # (bsz, seq_len, num_heads, head_dim), (1, 8192, 8, 128) - k = k.unsqueeze(0) # (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) - v = v.unsqueeze(0) # (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) - # cos, sin: (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) - cos = rotary_pos_emb_cos - sin = rotary_pos_emb_sin + # q: (bsz, seq_len, num_heads, head_dim), (1, 8192, 8, 128) + q = q.unflatten(-1, (self.num_heads, self.head_dim)).unsqueeze(0) + # k, v: (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0) + v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0) q, k = apply_rotary_pos_emb_sep( q, k, - cos, - sin, + rotary_pos_emb_cos, + rotary_pos_emb_sin, ) - q = q.squeeze(0) # (seq_len, num_heads, head_dim), (8192, 8, 128) - k = k.squeeze(0) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) - v = v.squeeze(0) # (seq_len, num_kv_heads, head_dim), (8192, 1, 128) - q = q.flatten(1, 2) # (seq_len, q_size), (8192, 1024) - k = k.flatten(1, 2) # (seq_len, kv_size), (8192, 128) - v = v.flatten(1, 2) # (seq_len, kv_size), (8192, 128) - - attn_output = self.attn(q, k, v) # (seq_len, q_size), (8192, 1024) + # q: (seq_len, num_heads, head_dim), (8192, 8*128) + q = q.squeeze(0).flatten(1, 2) + # k, v: (seq_len, num_kv_heads, head_dim), (8192, 1*128) + k = k.squeeze(0).flatten(1, 2) + v = v.squeeze(0).flatten(1, 2) + + # (seq_len, q_size), (8192, 1024) + attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output From 56ea6805bede0a2746de43557d0cd0908d896707 Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 7 Jan 2026 21:22:46 +0800 Subject: [PATCH 07/24] minor Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 249a45738580..5fb9e4406d71 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -196,11 +196,12 @@ def forward( rotary_pos_emb_sin: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) - # qkv: (8192, 1280) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # hidden_states: (seq_len, hidden_size), (8192, 2048) + # cos, sin: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) + # qkv: (8192, 1280) # q: (seq_len, q_size), (8192, 1024) - # k: (seq_len, kv_size), (8192, 128) - # v: (seq_len, kv_size), (8192, 128) + # k, v: (seq_len, kv_size), (8192, 128) # Add qk-norm q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) @@ -211,12 +212,7 @@ def forward( k_by_head = self.k_norm(k_by_head) k = k_by_head.view(k.shape) - # hidden_states: (seq_len, hidden_size), (8192, 2048) - # q: (seq_len, q_size), (8192, 1024) - # k, v: (seq_len, kv_size), (8192, 128) - # cos, sin: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) - - # apply rotary embedding sep head + # Apply rotary embedding sep head # q: (bsz, seq_len, num_heads, head_dim), (1, 8192, 8, 128) q = q.unflatten(-1, (self.num_heads, self.head_dim)).unsqueeze(0) # k, v: (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) From 86c301a6a48df9bf07dee3f23d75cf869b57de9d Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 7 Jan 2026 21:25:32 +0800 Subject: [PATCH 08/24] minor Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 5fb9e4406d71..249b46a01871 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -224,9 +224,9 @@ def forward( rotary_pos_emb_cos, rotary_pos_emb_sin, ) - # q: (seq_len, num_heads, head_dim), (8192, 8*128) + # q: (seq_len, q_size), (8192, 8*128) q = q.squeeze(0).flatten(1, 2) - # k, v: (seq_len, num_kv_heads, head_dim), (8192, 1*128) + # k, v: (seq_len, kv_size), (8192, 1*128) k = k.squeeze(0).flatten(1, 2) v = v.squeeze(0).flatten(1, 2) From 01d312cc743d245a21d032a081551672a2f9ba22 Mon Sep 17 00:00:00 2001 From: zxy Date: Thu, 8 Jan 2026 21:33:02 +0800 Subject: [PATCH 09/24] support group router Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 232 ++++++++++++++++++++- 1 file changed, 228 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 249b46a01871..65656295bb07 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -24,6 +24,7 @@ # limitations under the License. """Inference-only InternS1Pro model compatible with HuggingFace weights.""" +import functools import typing from collections.abc import Callable, Iterable from itertools import islice @@ -34,16 +35,22 @@ from vllm.attention.layer import Attention from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( + get_ep_group, get_pp_group, get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, ) from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -57,13 +64,13 @@ default_weight_loader, maybe_remap_kv_scale_name, ) +from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors +from .interfaces import MixtureOfExperts from .qwen3_moe import ( Qwen3MoeForCausalLM, - Qwen3MoeMLP, - Qwen3MoeSparseMoeBlock, ) from .qwen3_vl import ( Qwen3_VisionTransformer, @@ -72,7 +79,6 @@ Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, ) -from .qwen3_vl_moe import Qwen3VLMoeMixtureOfExperts from .utils import ( extract_layer_index, is_pp_missing_parameter, @@ -236,6 +242,182 @@ def forward( return output +class Qwen3MoeMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Qwen3MoeSparseMoeBlock(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + config = vllm_config.model_config.hf_text_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + self.tp_size = get_tensor_model_parallel_world_size() + + self.ep_group = get_ep_group().device_group + self.ep_rank = get_ep_group().rank_in_group + self.ep_size = self.ep_group.size() + self.n_routed_experts = config.num_experts + + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + + # Load balancing settings. + vllm_config = get_current_vllm_config() + eplb_config = vllm_config.parallel_config.eplb_config + self.enable_eplb = parallel_config.enable_eplb + + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + # For custom routing function + self.n_groups = getattr(config, "router_n_groups", -1) + + self.experts = FusedMoE( + num_experts=self.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=True, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + is_sequence_parallel=self.is_sequence_parallel, + routing_method_type=RoutingMethodType.Renormalize, + custom_routing_function=self._custom_routing_function, + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate", + ) + + @staticmethod + @functools.lru_cache + def get_group_offsets(n_groups: int, group_size: int, device: str): + group_offsets = (torch.arange(n_groups, device=device) * group_size).view( + 1, -1, 1 + ) # [1, n_groups, 1] + return group_offsets + + # TODO: zhouxinyu, use vllm routing functions + def _custom_routing_function( + self, + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> torch.Tensor: + routing_weights = torch.softmax(gating_output, dim=-1, dtype=torch.float32) + + if self.n_groups > 0: + assert routing_weights.shape[-1] % self.n_groups == 0, ( + f"{routing_weights.shape[-1]} cannot be divided by {self.n_groups}" + ) + per_group_top_k = topk // self.n_groups + group_size = routing_weights.shape[-1] // self.n_groups + group_offsets = self.get_group_offsets( + self.n_groups, group_size, routing_weights.device + ) + routing_weights = routing_weights.unflatten(-1, (self.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 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.dim() <= 2, ( + "Qwen3MoeSparseMoeBlock only supports 1D or 2D inputs" + ) + is_input_1d = hidden_states.dim() == 1 + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + if self.is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + if self.is_sequence_parallel: + final_hidden_states = tensor_model_parallel_all_gather( + final_hidden_states, 0 + ) + final_hidden_states = final_hidden_states[:num_tokens] + + # return to 1d if input is 1d + return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states + + class InternS1ProMoeDecoderLayer(nn.Module): def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -656,6 +838,48 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) +class Qwen3VLMoeMixtureOfExperts(MixtureOfExperts): + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.language_model.model.layers: + if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + def set_moe_parameters(self): + self.expert_weights = [] + + self.moe_layers = [] + example_moe = None + for layer in self.language_model.model.layers: + if hasattr(layer, "mlp") and isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + example_moe = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_moe is None: + raise RuntimeError("No Qwen3Moe layer found in the language_model.") + + # Set MoE hyperparameters + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + @MULTIMODAL_REGISTRY.register_processor( Qwen3VLMultiModalProcessor, info=InternS1ProProcessingInfo, From f73433763a620efb925e134218dbcbd85c69bcb8 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 9 Jan 2026 14:34:45 +0800 Subject: [PATCH 10/24] fix merge bug, rename Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 65656295bb07..8ad7f4b7d77c 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -242,7 +242,7 @@ def forward( return output -class Qwen3MoeMLP(nn.Module): +class InternS1ProMoeMLP(nn.Module): def __init__( self, hidden_size: int, @@ -281,7 +281,7 @@ def forward(self, x): return x -class Qwen3MoeSparseMoeBlock(nn.Module): +class InternS1ProMoeSparseMoeBlock(nn.Module): def __init__( self, vllm_config: VllmConfig, @@ -393,7 +393,7 @@ def _custom_routing_function( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: assert hidden_states.dim() <= 2, ( - "Qwen3MoeSparseMoeBlock only supports 1D or 2D inputs" + "InternS1ProMoeSparseMoeBlock only supports 1D or 2D inputs" ) is_input_1d = hidden_states.dim() == 1 num_tokens, hidden_dim = hidden_states.shape @@ -447,11 +447,11 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: if (layer_idx not in mlp_only_layers) and ( config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0 ): - self.mlp = Qwen3MoeSparseMoeBlock( + self.mlp = InternS1ProMoeSparseMoeBlock( vllm_config=vllm_config, prefix=f"{prefix}.mlp" ) else: - self.mlp = Qwen3MoeMLP( + self.mlp = InternS1ProMoeMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, @@ -614,6 +614,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) return FusedMoE.make_expert_params_mapping( + self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", @@ -849,7 +850,7 @@ def update_physical_experts_metadata( self.num_local_physical_experts = num_local_physical_experts self.num_redundant_experts = num_physical_experts - self.num_logical_experts for layer in self.language_model.model.layers: - if isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + if isinstance(layer.mlp, InternS1ProMoeSparseMoeBlock): moe = layer.mlp moe.n_local_physical_experts = num_local_physical_experts moe.n_physical_experts = num_physical_experts @@ -862,12 +863,14 @@ def set_moe_parameters(self): self.moe_layers = [] example_moe = None for layer in self.language_model.model.layers: - if hasattr(layer, "mlp") and isinstance(layer.mlp, Qwen3MoeSparseMoeBlock): + if hasattr(layer, "mlp") and isinstance( + layer.mlp, InternS1ProMoeSparseMoeBlock + ): example_moe = layer.mlp self.moe_layers.append(layer.mlp.experts) if example_moe is None: - raise RuntimeError("No Qwen3Moe layer found in the language_model.") + raise RuntimeError("No InternS1ProMoe layer found in the language_model.") # Set MoE hyperparameters self.num_moe_layers = len(self.moe_layers) From 42d2f3b7a8f373527b89175fed318ef6809e2cfe Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 12 Jan 2026 17:50:00 +0800 Subject: [PATCH 11/24] refactor Signed-off-by: zxy --- .../layers/rotary_embedding/__init__.py | 1 + .../layers/rotary_embedding/base.py | 21 +- .../layers/rotary_embedding/fope.py | 215 +++--- vllm/model_executor/models/interns1_pro.py | 661 +++++------------- vllm/model_executor/models/qwen3_moe.py | 10 +- vllm/model_executor/models/qwen3_vl_moe.py | 15 +- 6 files changed, 325 insertions(+), 598 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index 5db30834b480..20fd357e82e0 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -114,6 +114,7 @@ def get_rope( "fope_init_factor", ) } + extra_kwargs["init_cache"] = False rotary_emb = FourierRotaryEmbedding( head_size, rotary_dim, diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index bd82728ed15f..05b3109eed61 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -22,6 +22,7 @@ def __init__( base: float, is_neox_style: bool, dtype: torch.dtype, + init_cache: bool = True, ) -> None: super().__init__() self.head_size = head_size @@ -43,11 +44,12 @@ def __init__( if not hasattr(self, "use_flashinfer"): self.use_flashinfer = False - cache = self._compute_cos_sin_cache() - if not self.use_flashinfer: - cache = cache.to(dtype) - self.cos_sin_cache: torch.Tensor - self.register_buffer("cos_sin_cache", cache, persistent=False) + if init_cache: + cache = self._compute_cos_sin_cache() + if not self.use_flashinfer: + cache = cache.to(dtype) + self.cos_sin_cache: torch.Tensor + self.register_buffer("cos_sin_cache", cache, persistent=False) self.is_rocm_triton_rotary_embed_enabled = ( rocm_aiter_ops.is_triton_rotary_embed_enabled() ) @@ -105,9 +107,16 @@ def __init__( base: float, is_neox_style: bool, dtype: torch.dtype, + init_cache: bool = True, ) -> None: super().__init__( - head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + head_size, + rotary_dim, + max_position_embeddings, + base, + is_neox_style, + dtype, + init_cache, ) @staticmethod diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index 071e374be56b..671f13b7f86d 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -3,9 +3,9 @@ import torch import torch.nn.functional as F from torch import nn -from torch.nn.parameter import Parameter from .base import RotaryEmbedding +from .common import rotate_neox class FourierRotaryEmbedding(RotaryEmbedding): @@ -17,6 +17,7 @@ def __init__( base: float, is_neox_style: bool, dtype: torch.dtype, + init_cache: bool, # extra parameters for FoPE num_key_value_heads: int, num_inv_freq: int, @@ -30,10 +31,21 @@ def __init__( self.fope_init_factor = fope_init_factor super().__init__( - head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + head_size, + rotary_dim, + max_position_embeddings, + base, + is_neox_style, + dtype, + init_cache, + ) + + # setup buffers and parameters + self.inv_freq: torch.Tensor + self.register_buffer( + "inv_freq", self._compute_inv_freq(self.base), persistent=False ) - # setup parameters self.input_dim = self.inv_freq.shape[-1] self.output_dim = self.inv_freq.shape[-1] self.cos_coef = nn.Parameter( @@ -44,8 +56,13 @@ def __init__( torch.empty(num_key_value_heads, self.input_dim, self.output_dim), requires_grad=False, ) - self.cos_coef.weight_loader = self.fope_coef_weight_loader - self.sin_coef.weight_loader = self.fope_coef_weight_loader + + self.cos_sin_cache: torch.Tensor + cache = self._compute_cos_sin_cache().to(dtype) + self.register_buffer("cos_sin_cache", cache, persistent=False) + + # update cache in the first forward, where sin/cos_coef weights are ready + self.update_cache = True def _compute_inv_freq(self, base: float) -> torch.Tensor: """Compute the inverse frequency.""" @@ -58,122 +75,104 @@ def _compute_inv_freq(self, base: float) -> torch.Tensor: inv_freq_idx_selected = torch.ones_like(inv_freq, dtype=torch.bool) if self.num_inv_freq is not None: - num_inv_freq = self.num_inv_freq - inv_freq_idx_selected[num_inv_freq:] = False + inv_freq_idx_selected[self.num_inv_freq :] = False else: inv_freq_idx_selected = inv_freq > ( 2.0 * torch.pi / self.max_position_embeddings ) - num_inv_freq = inv_freq_idx_selected.sum().item() 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.""" - self.inv_freq = self._compute_inv_freq(self.base) - # TODO: zhouxinyu, implement FoPE cos/sin cache computation - return torch.zeros(1) - - def apply_rotary_emb(self): - """Customized apply_rotary_emb function for FoPE.""" - pass - - def get_step_eye(self, _param): - import math - - _step_eye = torch.zeros_like(_param) - - step = math.ceil(self.input_dim / self.output_dim) - for i in range(self.output_dim): - if i * step < self.input_dim: - _step_eye[..., i * step, i] = 1.0 - - return _step_eye - - def forward_native(self, x: torch.Tensor, positions: torch.Tensor): - # expand x, positions to additional batch size dim - if x.dim() == 2: - # (seq_len, hidden_size) -> (bsz, seq_len, hidden_size) - x = x.unsqueeze(0) - if positions.dim() == 1: - # (seq_len) -> (bsz, seq_len) - positions = positions.unsqueeze(0) - - # Core RoPE block - inv_freq_expanded = ( - self.inv_freq[None, :, None].float().expand(positions.shape[0], -1, 1) - ) # (40) -> (1, 40, 1) -> (bsz, 40, 1) - position_ids_expanded = positions[ - :, None, : - ].float() # (bsz, seq_len) -> (bsz, 1, seq_len) - # Force float32 (see https://github.com/huggingface/transformers/pull/29285) - device_type = x.device.type - batch_size, seq_len, hidden_size = x.shape - with torch.autocast(device_type=device_type, enabled=False): - freqs = ( - inv_freq_expanded.float() @ position_ids_expanded.float() - ).transpose(1, 2) - if self.fope_sep_head: - pos_cos = ( - freqs.cos() - .unsqueeze(1) - .expand(batch_size, self.num_key_value_heads, seq_len, -1) - ) - pos_sin = ( - freqs.sin() - .unsqueeze(1) - .expand(batch_size, self.num_key_value_heads, seq_len, -1) - ) - else: - pos_cos = freqs.cos() - pos_sin = freqs.sin() - - if self.fope_sep_head: - # (1, 1, 8192, 40) x (1, 40, 40) -> (1, 8192, 1, 40) - sin = torch.einsum("bhtD, hDd -> bthd", pos_sin, self.sin_coef.float()) - cos = torch.einsum("bhtD, hDd -> bthd", pos_cos, self.cos_coef.float()) - else: - sin = torch.einsum("btD, Dd -> btd", pos_sin, self.sin_coef.float()) - cos = torch.einsum("btD, Dd -> btd", 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, - ) + device = self.inv_freq.device + t = torch.arange(self.max_position_embeddings, dtype=torch.float, device=device) - sin = torch.cat((sin, sin), dim=-1) - cos = torch.cat((cos, cos), dim=-1) - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - def forward_cuda(self, x: torch.Tensor, positions: torch.Tensor): - # TODO, zhouxinyu, implement FoPE cuda forward computation - return self.forward_native(x, positions) + freqs = torch.einsum("j,i -> ji", t, self.inv_freq) + if self.fope_sep_head: + pos_cos = freqs.cos().unsqueeze(0).expand(self.num_key_value_heads, -1, -1) + pos_sin = freqs.sin().unsqueeze(0).expand(self.num_key_value_heads, -1, -1) + else: + pos_cos = freqs.cos() + pos_sin = freqs.sin() - def fope_coef_weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): - from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, + 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, ) - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - num_key_value_heads = loaded_weight.size(0) + sin = torch.cat((sin, sin), dim=-1) + cos = torch.cat((cos, cos), dim=-1) - if num_key_value_heads < world_size: - n_replicate = world_size // num_key_value_heads - world_size = num_key_value_heads - rank = rank // n_replicate + # cache: (max_position_embeddings, num_kv_heads, kv_size * 2) + cache = torch.cat((cos, sin), dim=-1) + return cache - loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] - param.copy_(loaded_weight) + def forward_native( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # update cos/sin cache in the first forward + if self.update_cache: + cache = self._compute_cos_sin_cache().to(self.dtype) + self.cos_sin_cache.copy_(cache) + self.update_cache = False + + positions = positions.flatten() + cos_sin = self.cos_sin_cache.index_select(0, positions) + cos, sin = cos_sin.chunk(2, dim=-1) + + # apply rotary embedding + # query: (seq_len, num_heads, head_size) + # key: (seq_len, num_kv_heads, head_size) + query = query.unflatten(-1, (-1, self.head_size)) + assert key is not None, "Key tensor is required for FoPE." + key = key.unflatten(-1, (-1, self.head_size)) + + 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.view(seq_len, -1, query.size(-1)) + key = key.view(seq_len, -1, key.size(-1)) + + # native implementation of apply rope for neox style + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + query = (query * cos) + (rotate_neox(query) * sin) + key = (key * cos) + (rotate_neox(key) * sin) + + if need_reshape: + query = query.view(query_shape) + key = key.view(key_shape) + + return query, key diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 8ad7f4b7d77c..801c7138c786 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -25,20 +25,18 @@ """Inference-only InternS1Pro model compatible with HuggingFace weights.""" import functools -import typing -from collections.abc import Callable, Iterable -from itertools import islice +from collections.abc import Iterable +from typing import Any import torch from torch import nn from transformers import AutoProcessor, PretrainedConfig from vllm.attention.layer import Attention -from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( get_ep_group, - get_pp_group, + get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, ) @@ -58,15 +56,12 @@ from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, - VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, - maybe_remap_kv_scale_name, ) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.sequence import IntermediateTensors from .interfaces import MixtureOfExperts from .qwen3_moe import ( @@ -79,11 +74,9 @@ Qwen3VLMultiModalProcessor, Qwen3VLProcessingInfo, ) +from .qwen3_vl_moe import Qwen3MoeLLMModel from .utils import ( extract_layer_index, - is_pp_missing_parameter, - make_empty_intermediate_tensors_factory, - make_layers, maybe_prefix, ) @@ -102,146 +95,6 @@ def get_hf_processor(self, **kwargs: object) -> AutoProcessor: ) -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand( - batch, num_key_value_heads, n_rep, slen, head_dim - ) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def apply_rotary_pos_emb_sep(q, k, cos, sin, unsqueeze_dim=1): - num_groups = int(q.shape[unsqueeze_dim] // cos.shape[unsqueeze_dim]) - cos_rep = repeat_kv(cos, num_groups) - sin_rep = repeat_kv(sin, num_groups) - q_embed = (q * cos_rep) + (rotate_half(q) * sin_rep) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - -class InternS1ProMoeAttention(nn.Module): - def __init__( - self, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int | None = None, - rms_norm_eps: float = 1e-06, - qkv_bias: bool = False, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = head_dim or (hidden_size // self.total_num_heads) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=qkv_bias, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - - self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - - # TODO: zhouxinyu, use vllm ApplyRotaryEmb - # self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) - - def forward( - self, - hidden_states: torch.Tensor, - rotary_pos_emb_cos: torch.Tensor, - rotary_pos_emb_sin: torch.Tensor, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - # hidden_states: (seq_len, hidden_size), (8192, 2048) - # cos, sin: (bsz, seq_len, num_kv_heads, kv_size), (1, 8192, 1, 128) - # qkv: (8192, 1280) - # q: (seq_len, q_size), (8192, 1024) - # k, v: (seq_len, kv_size), (8192, 128) - - # Add qk-norm - q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) - q_by_head = self.q_norm(q_by_head) - q = q_by_head.view(q.shape) - - k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) - k_by_head = self.k_norm(k_by_head) - k = k_by_head.view(k.shape) - - # Apply rotary embedding sep head - # q: (bsz, seq_len, num_heads, head_dim), (1, 8192, 8, 128) - q = q.unflatten(-1, (self.num_heads, self.head_dim)).unsqueeze(0) - # k, v: (bsz, seq_len, num_kv_heads, head_dim), (1, 8192, 1, 128) - k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0) - v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0) - q, k = apply_rotary_pos_emb_sep( - q, - k, - rotary_pos_emb_cos, - rotary_pos_emb_sin, - ) - # q: (seq_len, q_size), (8192, 8*128) - q = q.squeeze(0).flatten(1, 2) - # k, v: (seq_len, kv_size), (8192, 1*128) - k = k.squeeze(0).flatten(1, 2) - v = v.squeeze(0).flatten(1, 2) - - # (seq_len, q_size), (8192, 1024) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - class InternS1ProMoeMLP(nn.Module): def __init__( self, @@ -418,6 +271,111 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states +class InternS1ProMoeAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_parameters: dict[str, Any], + max_position_embeddings: int = 32768, + head_dim: int | None = None, + rms_norm_eps: float = 1e-06, + qkv_bias: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + dual_chunk_attention_config: dict[str, Any] | None = None, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim or (hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.dual_chunk_attention_config = dual_chunk_attention_config + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + rope_parameters["num_key_value_heads"] = self.num_kv_heads + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + **{ + "layer_idx": extract_layer_index(prefix), + "dual_chunk_attention_config": dual_chunk_attention_config, + } + if dual_chunk_attention_config + else {}, + ) + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # Add qk-norm + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) + q_by_head = self.q_norm(q_by_head) + q = q_by_head.view(q.shape) + + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) + k_by_head = self.k_norm(k_by_head) + k = k_by_head.view(k.shape) + q, k = self.rotary_emb.forward_native(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + class InternS1ProMoeDecoderLayer(nn.Module): def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -427,16 +385,38 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: quant_config = vllm_config.quant_config self.hidden_size = config.hidden_size + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + + # update rope related parameters + rope_scaling = config.rope_scaling + 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) + fope_init_factor = rope_scaling.get("fope_init_factor", None) + fope_sep_head = rope_scaling.get("fope_sep_head", None) + num_inv_freq = rope_scaling.get("num_inv_freq", None) + + config.rope_parameters["use_fope"] = use_fope + config.rope_parameters["fope_init_factor"] = fope_init_factor + config.rope_parameters["fope_sep_head"] = fope_sep_head + config.rope_parameters["num_inv_freq"] = num_inv_freq + + assert use_fope, "should use FOPE for InternS1Pro model" self.self_attn = InternS1ProMoeAttention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, + rope_parameters=config.rope_parameters, + max_position_embeddings=max_position_embeddings, rms_norm_eps=config.rms_norm_eps, qkv_bias=getattr(config, "attention_bias", False), head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", + dual_chunk_attention_config=dual_chunk_attention_config, ) # `mlp_only_layers` in the config. @@ -465,9 +445,8 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: def forward( self, + positions: torch.Tensor, hidden_states: torch.Tensor, - rotary_pos_emb_cos: torch.Tensor, - rotary_pos_emb_sin: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention @@ -477,9 +456,8 @@ def forward( else: hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn( + positions=positions, hidden_states=hidden_states, - rotary_pos_emb_cos=rotary_pos_emb_cos, - rotary_pos_emb_sin=rotary_pos_emb_sin, ) # Fully Connected @@ -488,333 +466,19 @@ def forward( return hidden_states, residual -@support_torch_compile( - dynamic_arg_dims={ - "input_ids": 0, - "positions": -1, - "intermediate_tensors": 0, - "inputs_embeds": 0, - } -) -class InternS1ProMoeLLMModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_text_config - quant_config = vllm_config.quant_config - parallel_config = vllm_config.parallel_config - eplb_config = parallel_config.eplb_config - self.num_redundant_experts = eplb_config.num_redundant_experts - - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - self.config = config - self.quant_config = quant_config - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - - # build rotary embedding - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - tp_size = get_tensor_model_parallel_world_size() - num_key_value_heads = max(1, config.num_key_value_heads // tp_size) - dual_chunk_attention_config = getattr( - config, "dual_chunk_attention_config", None - ) - head_dim = getattr(config, "head_dim", None) - if head_dim is None: - head_dim = config.hidden_size // config.num_attention_heads - - rope_scaling = config.rope_scaling - fope_init_factor = rope_scaling.get("fope_init_factor", None) - fope_sep_head = rope_scaling.get("fope_sep_head", None) - num_inv_freq = rope_scaling.get("num_inv_freq", None) - - 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) - rope_parameters = { - "rope_theta": config.rope_theta, - "num_key_value_heads": num_key_value_heads, - "use_fope": use_fope, - "fope_init_factor": fope_init_factor, - "fope_sep_head": fope_sep_head, - "num_inv_freq": num_inv_freq, - } - self.rotary_emb = get_rope( - head_size=head_dim, - max_position=max_position_embeddings, - rope_parameters=rope_parameters, - dual_chunk_attention_config=dual_chunk_attention_config, - ) - - # build layers - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: InternS1ProMoeDecoderLayer( - vllm_config=vllm_config, prefix=prefix - ), - prefix=f"{prefix}.layers", - ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - # Track layers for auxiliary hidden state outputs (EAGLE3) - self.aux_hidden_state_layers: tuple[int, ...] = () - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - deepstack_input_embeds: IntermediateTensors | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - # rotary embedding - rotary_pos_emb_cos, rotary_pos_emb_sin = self.rotary_emb( - hidden_states, positions - ) - - # decoder layers - for layer_idx, layer in islice( - enumerate(self.layers), self.start_layer, self.end_layer - ): - hidden_states, residual = layer( - hidden_states=hidden_states, - rotary_pos_emb_cos=rotary_pos_emb_cos, - rotary_pos_emb_sin=rotary_pos_emb_sin, - residual=residual, - ) - - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - return FusedMoE.make_expert_params_mapping( - self, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - num_redundant_experts=self.num_redundant_experts, - ) - - def load_fused_expert_weights( +class InternS1ProMoeLLMModel(Qwen3MoeLLMModel): + def __init__( self, - name: str, - params_dict: dict, - loaded_weight: torch.Tensor, - shard_id: str, - num_experts: int, - ) -> bool: - param = params_dict[name] - weight_loader = typing.cast(Callable[..., bool], param.weight_loader) - loaded_local_expert = False - for expert_id in range(num_experts): - curr_expert_weight = loaded_weight[expert_id] - success = weight_loader( - param, - curr_expert_weight, - name, - shard_id, - expert_id, - return_success=True, - ) - if success: - loaded_local_expert = True - - return loaded_local_expert - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - # Skip loading extra parameters for GPTQ/modelopt models. - ignore_suffixes = ( - ".bias", - "_bias", - ".k_scale", - "_k_scale", - ".v_scale", - "_v_scale", - ".weight_scale", - "_weight_scale", - ".input_scale", - "_input_scale", + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[torch.nn.Module] = InternS1ProMoeDecoderLayer, + ): + super().__init__( + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=decoder_layer_type, ) - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - expert_params_mapping = self.get_expert_mapping() - is_fused_expert = False - fused_expert_params_mapping = [ - ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), - ("experts.w2_weight", "experts.down_proj", 0, "w2"), - ] - num_experts = self.config.num_experts - for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: - if "experts.gate_up_proj" in name or "experts.down_proj" in name: - is_fused_expert = True - expert_params_mapping = fused_expert_params_mapping - - # Skip non-stacked layers and experts (experts handled below). - if weight_name not in name: - continue - # We have mlp.experts[0].gate_proj in the checkpoint. - # Since we handle the experts below in expert_params_mapping, - # we need to skip here BEFORE we update the name, otherwise - # name will be updated to mlp.experts[0].gate_up_proj, which - # will then be updated below in expert_params_mapping - # for mlp.experts[0].gate_gate_up_proj, which breaks load. - if "mlp.experts" in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra parameters for GPTQ/modelopt models. - if name.endswith(ignore_suffixes) and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - if name.endswith("scale"): - # Remapping the name of FP8 kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - if weight_loader == default_weight_loader: - weight_loader(param, loaded_weight) - else: - weight_loader(param, loaded_weight, shard_id) - break - else: - is_expert_weight = False - for mapping in expert_params_mapping: - param_name, weight_name, expert_id, shard_id = mapping - if weight_name not in name: - continue - # Anyway, this is an expert weight and should not be - # attempted to load as other weights later - is_expert_weight = True - name_mapped = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name_mapped, self): - continue - if is_fused_expert: - loaded_weight = loaded_weight.transpose(-1, -2) # no bias - if "experts.gate_up_proj" in name: - loaded_weight = loaded_weight.chunk(2, dim=-2) - success_w1 = self.load_fused_expert_weights( - name_mapped, - params_dict, - loaded_weight[0], - "w1", - num_experts, - ) - success_w3 = self.load_fused_expert_weights( - name_mapped, - params_dict, - loaded_weight[1], - "w3", - num_experts, - ) - success = success_w1 and success_w3 - else: - # down_proj - success = self.load_fused_expert_weights( - name_mapped, - params_dict, - loaded_weight, - shard_id, - num_experts, - ) - else: - # Skip loading extra parameters for GPTQ/modelopt models - if ( - name_mapped.endswith(ignore_suffixes) - and name_mapped not in params_dict - ): - continue - param = params_dict[name_mapped] - # We should ask the weight loader to return success or - # not here since otherwise we may skip experts with - # other available replicas. - weight_loader = typing.cast( - Callable[..., bool], param.weight_loader - ) - success = weight_loader( - param, - loaded_weight, - name_mapped, - shard_id=shard_id, - expert_id=expert_id, - return_success=True, - ) - if success: - name = name_mapped - break - else: - if is_expert_weight: - # We've checked that this is an expert weight - # However it's not mapped locally to this rank - # So we simply skip it - continue - # Skip loading extra parameters for GPTQ/modelopt models. - if name.endswith(ignore_suffixes) and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - # Remapping the name of FP8 kv-scale. - if name.endswith("kv_scale"): - remapped_kv_scale_name = name.replace( - ".kv_scale", ".attn.kv_scale" - ) - if remapped_kv_scale_name not in params_dict: - logger.warning_once( - "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.", # noqa: E501 - name, - remapped_kv_scale_name, - ) - continue - else: - name = remapped_kv_scale_name - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM): @@ -951,3 +615,40 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Set MoE hyperparameters self.set_moe_parameters() + + def _load_fope_weights(self, name: str, loaded_weight: torch.Tensor, params_dict): + """load fope weights""" + world_size = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + num_key_value_heads = loaded_weight.size(0) + + if num_key_value_heads < world_size: + n_replicate = world_size // num_key_value_heads + world_size = num_key_value_heads + rank = rank // n_replicate + + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] + + # rotary_emb is shared across all layers + param_name = name.replace(".rotary_emb.", ".layers.0.self_attn.rotary_emb.") + assert param_name in params_dict, f"{param_name} not in model parameters" + 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"language_model.model.") + self._load_fope_weights(name, loaded_weight, params_dict) + else: + other_weights[name] = loaded_weight + + super().load_weights(other_weights.items()) diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index f2f3546047aa..b100321e3c68 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -390,7 +390,13 @@ def forward( @support_torch_compile class Qwen3MoeModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[torch.nn.Module] = Qwen3MoeDecoderLayer, + ): super().__init__() config = vllm_config.model_config.hf_text_config @@ -411,7 +417,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, - lambda prefix: Qwen3MoeDecoderLayer(vllm_config=vllm_config, prefix=prefix), + lambda prefix: decoder_layer_type(vllm_config=vllm_config, prefix=prefix), prefix=f"{prefix}.layers", ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 3186804488e5..e360fb97833b 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -48,6 +48,7 @@ from .interfaces import MixtureOfExperts from .qwen3_moe import ( + Qwen3MoeDecoderLayer, Qwen3MoeForCausalLM, Qwen3MoeModel, Qwen3MoeSparseMoeBlock, @@ -82,8 +83,18 @@ def get_hf_config(self): } ) class Qwen3MoeLLMModel(Qwen3MoeModel): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__(vllm_config=vllm_config, prefix=prefix) + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[torch.nn.Module] = Qwen3MoeDecoderLayer, + ): + super().__init__( + vllm_config=vllm_config, + prefix=prefix, + decoder_layer_type=decoder_layer_type, + ) if not get_pp_group().is_first_rank: assert self.start_layer >= len( vllm_config.model_config.hf_config.vision_config.deepstack_visual_indexes From 3d91fbef3fe80ffa90ce51bb9f4f997f48329563 Mon Sep 17 00:00:00 2001 From: zxy Date: Wed, 14 Jan 2026 17:22:40 +0800 Subject: [PATCH 12/24] fix gate quant error Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 801c7138c786..c8142598ac6d 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -199,7 +199,6 @@ def __init__( config.hidden_size, config.num_experts, bias=False, - quant_config=quant_config, prefix=f"{prefix}.gate", ) From bc4f14546fd12fa27181e157a6040ae206b672be Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 16 Jan 2026 13:06:24 +0800 Subject: [PATCH 13/24] fix visual quant error Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index c8142598ac6d..2b022cde0dd6 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -566,7 +566,6 @@ class InternS1ProForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super(Qwen3VLForConditionalGeneration, self).__init__() config: PretrainedConfig = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config self.config = config @@ -585,7 +584,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.visual = Qwen3_VisionTransformer( config.vision_config, norm_eps=getattr(config, "rms_norm_eps", 1e-6), - quant_config=quant_config, multimodal_config=multimodal_config, prefix=maybe_prefix(prefix, "visual"), ) From 44cfb07e0b814332c1e8ccc6d8ef280d1ab10bbd Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sat, 17 Jan 2026 23:36:40 +0800 Subject: [PATCH 14/24] clean fope code Signed-off-by: Isotr0py --- .../layers/rotary_embedding/fope.py | 21 +++++++ vllm/model_executor/models/interns1_pro.py | 59 +++++++------------ 2 files changed, 42 insertions(+), 38 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index 671f13b7f86d..ccbacbdc44f3 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -4,6 +4,11 @@ import torch.nn.functional as F from torch import nn +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) + from .base import RotaryEmbedding from .common import rotate_neox @@ -56,6 +61,8 @@ def __init__( torch.empty(num_key_value_heads, self.input_dim, self.output_dim), requires_grad=False, ) + self.sin_coef.weight_loader = self.weight_loader + self.cos_coef.weight_loader = self.weight_loader self.cos_sin_cache: torch.Tensor cache = self._compute_cos_sin_cache().to(dtype) @@ -176,3 +183,17 @@ def forward_native( key = key.view(key_shape) return query, key + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + """load fope weights""" + world_size = get_tensor_model_parallel_world_size() + rank = get_tensor_model_parallel_rank() + num_key_value_heads = loaded_weight.size(0) + + if num_key_value_heads < world_size: + n_replicate = world_size // num_key_value_heads + world_size = num_key_value_heads + rank = rank // n_replicate + + loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] + param.data.copy_(loaded_weight) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 2b022cde0dd6..e415152e817d 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -36,7 +36,6 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( get_ep_group, - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, ) @@ -57,9 +56,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, ) -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, -) from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.multimodal import MULTIMODAL_REGISTRY @@ -76,6 +72,8 @@ ) from .qwen3_vl_moe import Qwen3MoeLLMModel from .utils import ( + AutoWeightsLoader, + WeightsMapper, extract_layer_index, maybe_prefix, ) @@ -563,6 +561,20 @@ class InternS1ProForConditionalGeneration( ], } + # To ensure correct weight loading and mapping. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.visual.": "visual.", + "lm_head.": "language_model.lm_head.", + "model.language_model.": "language_model.model.", + }, + orig_to_new_suffix={ + # Handle FOPE rotary embeddings + ".rotary_emb.sin_coef": ".layers.0.self_attn.rotary_emb.sin_coef", + ".rotary_emb.cos_coef": ".layers.0.self_attn.rotary_emb.cos_coef", + }, + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super(Qwen3VLForConditionalGeneration, self).__init__() config: PretrainedConfig = vllm_config.model_config.hf_config @@ -613,39 +625,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Set MoE hyperparameters self.set_moe_parameters() - def _load_fope_weights(self, name: str, loaded_weight: torch.Tensor, params_dict): - """load fope weights""" - world_size = get_tensor_model_parallel_world_size() - rank = get_tensor_model_parallel_rank() - num_key_value_heads = loaded_weight.size(0) - - if num_key_value_heads < world_size: - n_replicate = world_size // num_key_value_heads - world_size = num_key_value_heads - rank = rank // n_replicate - - loaded_weight = loaded_weight.chunk(world_size, dim=0)[rank] - - # rotary_emb is shared across all layers - param_name = name.replace(".rotary_emb.", ".layers.0.self_attn.rotary_emb.") - assert param_name in params_dict, f"{param_name} not in model parameters" - 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"language_model.model.") - self._load_fope_weights(name, loaded_weight, params_dict) - else: - other_weights[name] = loaded_weight - - super().load_weights(other_weights.items()) + skip_prefixes = ["model.time_series."] + if self.visual is None: + skip_prefixes.append("visual.") + loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) From 56ec59696db586c72b743590686fd5a257f24f63 Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 16 Jan 2026 17:51:41 +0800 Subject: [PATCH 15/24] remove unnecessary get config Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index e415152e817d..dfcb716e90fc 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -33,7 +33,7 @@ from transformers import AutoProcessor, PretrainedConfig from vllm.attention.layer import Attention -from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.config import CacheConfig, VllmConfig from vllm.distributed import ( get_ep_group, get_tensor_model_parallel_world_size, @@ -160,7 +160,6 @@ def __init__( ) # Load balancing settings. - vllm_config = get_current_vllm_config() eplb_config = vllm_config.parallel_config.eplb_config self.enable_eplb = parallel_config.enable_eplb From 5fdad3ba30fc7831d35bc6bbc3373ac0b4da0df8 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sat, 24 Jan 2026 15:52:42 +0800 Subject: [PATCH 16/24] fix video example Signed-off-by: Isotr0py --- examples/offline_inference/vision_language.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index 95de4348f19d..75cf48603bc0 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -744,9 +744,9 @@ def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData: ) if modality == "image": - placeholder = "<|image_pad|>" + placeholder = "<|vision_start|><|image_pad|><|vision_end|>" elif modality == "video": - placeholder = "<|video_pad|>" + placeholder = "<|vision_start|><|video_pad|><|vision_end|>" tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) messages = [ From 95f9f434981c91c06719b59933d4dd1014dbf400 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 10:59:00 +0800 Subject: [PATCH 17/24] add docs Signed-off-by: zxy --- docs/models/supported_models.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 2fa05c1c37bc..a60506ea86b7 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -677,6 +677,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | | | `IsaacForConditionalGeneration` | Isaac | T + I+ | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ | | `InternS1ForConditionalGeneration` | Intern-S1 | T + IE+ + VE+ | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ | +| `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + IE+ + VE+ | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ | | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ | | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ | From 72a2e8c4937ef41e2376b2690e68d0bf77c93ede Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 11:29:08 +0800 Subject: [PATCH 18/24] change example model name Signed-off-by: zxy --- examples/offline_inference/vision_language.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index 75cf48603bc0..f32e6da2132f 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -731,7 +731,7 @@ def run_interns1(questions: list[str], modality: str) -> ModelRequestData: # Intern-S1-Pro def run_interns1_pro(questions: list[str], modality: str) -> ModelRequestData: - model_name = "internlm-private/InternS1_1_1T_A22_1217" + model_name = "internlm/Intern-S1-Pro" engine_args = EngineArgs( model=model_name, From ce85d10b61c2c53360b78c1ba160008e2ed01362 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Tue, 3 Feb 2026 13:35:27 +0800 Subject: [PATCH 19/24] update registry Signed-off-by: Isotr0py --- tests/models/registry.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/models/registry.py b/tests/models/registry.py index ebff16cecdfb..8a40a11b507e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -712,6 +712,11 @@ def check_available_online( "InternS1ForConditionalGeneration": _HfExamplesInfo( "internlm/Intern-S1", trust_remote_code=True ), + "InternS1ProForConditionalGeneration": _HfExamplesInfo( + "internlm/Intern-S1-Pro", + trust_remote_code=True, + is_available_online=False, + ), "InternVLChatModel": _HfExamplesInfo( "OpenGVLab/InternVL2-1B", extras={ From 5dc600fe7f3d7dea7605f12466fcadcb39f22067 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 14:04:53 +0800 Subject: [PATCH 20/24] fix according to review comments Signed-off-by: zxy --- vllm/model_executor/models/interns1_pro.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index dfcb716e90fc..60c92cddab3f 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -479,7 +479,7 @@ def __init__( class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super(Qwen3MoeForCausalLM, self).__init__() + super().__init__() self.config = vllm_config.model_config.hf_config.text_config self.quant_config = vllm_config.quant_config self.model = InternS1ProMoeLLMModel( @@ -575,7 +575,7 @@ class InternS1ProForConditionalGeneration( ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super(Qwen3VLForConditionalGeneration, self).__init__() + super().__init__() config: PretrainedConfig = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config From 633d0d4dfdfc51416022cb96ae7312712dcc1840 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 14:22:56 +0800 Subject: [PATCH 21/24] explicit kwargs Signed-off-by: zxy --- .../model_executor/layers/rotary_embedding/base.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index 05b3109eed61..d92aeb9bc733 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -110,13 +110,13 @@ def __init__( init_cache: bool = True, ) -> None: super().__init__( - head_size, - rotary_dim, - max_position_embeddings, - base, - is_neox_style, - dtype, - init_cache, + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + init_cache=init_cache, ) @staticmethod From 855dee7fc927f1a60e913b80c21c2771c64f9ed7 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 14:33:32 +0800 Subject: [PATCH 22/24] explicit kwargs in fope Signed-off-by: zxy --- .../model_executor/layers/rotary_embedding/fope.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/fope.py b/vllm/model_executor/layers/rotary_embedding/fope.py index ccbacbdc44f3..4c8a7bcbfa1d 100644 --- a/vllm/model_executor/layers/rotary_embedding/fope.py +++ b/vllm/model_executor/layers/rotary_embedding/fope.py @@ -36,13 +36,13 @@ def __init__( self.fope_init_factor = fope_init_factor super().__init__( - head_size, - rotary_dim, - max_position_embeddings, - base, - is_neox_style, - dtype, - init_cache, + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=max_position_embeddings, + base=base, + is_neox_style=is_neox_style, + dtype=dtype, + init_cache=init_cache, ) # setup buffers and parameters From 284e119e97fa052a5a5ed782b3a42dbaabbf1f36 Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 16:34:02 +0800 Subject: [PATCH 23/24] add min transformers version --- tests/models/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/models/registry.py b/tests/models/registry.py index 2c052e663f8e..24434ce248c6 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -758,6 +758,7 @@ def check_available_online( "InternS1ProForConditionalGeneration": _HfExamplesInfo( "internlm/Intern-S1-Pro", trust_remote_code=True, + min_transformers_version="4.57", is_available_online=False, ), "InternVLChatModel": _HfExamplesInfo( From d2da143a863b624b1727b98c5cf31199e0cb120c Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 3 Feb 2026 16:37:39 +0800 Subject: [PATCH 24/24] fix version --- tests/models/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 24434ce248c6..c38637c1c675 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -758,7 +758,7 @@ def check_available_online( "InternS1ProForConditionalGeneration": _HfExamplesInfo( "internlm/Intern-S1-Pro", trust_remote_code=True, - min_transformers_version="4.57", + min_transformers_version="5.0.0", is_available_online=False, ), "InternVLChatModel": _HfExamplesInfo(