# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen3Next MTP model.""" from collections.abc import Iterable import torch from torch import nn from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from vllm.model_executor.models.qwen3_next import ( Qwen3NextDecoderLayer, Qwen3NextModel, Qwen3NextRMSNorm, QwenNextMixtureOfExperts, ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from .utils import ( AutoWeightsLoader, WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix, ) logger = init_logger(__name__) KVCache = tuple[torch.Tensor, torch.Tensor] @support_torch_compile class Qwen3NextMultiTokenPredictor(nn.Module): hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() model_config = vllm_config.model_config quant_config = vllm_config.quant_config config: Qwen3NextConfig = model_config.hf_config self.config = config self.vocab_size = config.vocab_size self.mtp_start_layer_idx = config.num_hidden_layers self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1) self.embed_tokens = VocabParallelEmbedding( self.vocab_size, config.hidden_size, ) # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is # missing from the checkpoint quant exclude list (its `ignore` glob # does not cover `mtp.fc`). Force unquantized to match the weights, # mirroring the Qwen3.5 MTP handling (PR #38832). fc_quant = ( None if (quant_config and quant_config.get_name() == "modelopt_fp4") else quant_config ) self.fc = ColumnParallelLinear( self.config.hidden_size * 2, self.config.hidden_size, gather_output=True, bias=False, return_bias=False, quant_config=fc_quant, prefix=f"{prefix}.fc", ) self.layers = torch.nn.ModuleList( Qwen3NextDecoderLayer( vllm_config, layer_type="full_attention", prefix=f"{prefix}.layers.{idx}", ) for idx in range(self.num_mtp_layers) ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) self.norm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.pre_fc_norm_hidden = Qwen3NextRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.pre_fc_norm_embedding = Qwen3NextRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) 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, hidden_states: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, spec_step_idx: int = 0, ) -> torch.Tensor: if get_pp_group().is_first_rank: if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) assert hidden_states.shape[-1] == inputs_embeds.shape[-1] inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds) hidden_states = self.pre_fc_norm_hidden(hidden_states) hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1) hidden_states = self.fc(hidden_states) residual = None else: assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] current_step_idx = spec_step_idx % self.num_mtp_layers hidden_states, residual = self.layers[current_step_idx]( positions=positions, hidden_states=hidden_states, 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 load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: mapper = self.hf_to_vllm_mapper if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled(): # AITER fused-shared-experts: route the shared_expert checkpoint # weights into the extra fused expert slot. num_routed = self.config.num_experts mapper = mapper | WeightsMapper( orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."} ) loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=mapper) @support_torch_compile class Qwen3NextMTP(nn.Module, QwenNextMixtureOfExperts): packed_modules_mapping = { "qkv_proj": [ "q_proj", "k_proj", "v_proj", ], "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.vllm_config = vllm_config cache_config = vllm_config.cache_config if cache_config.mamba_cache_mode == "all": raise NotImplementedError( "Qwen3NextMTP currently does not support 'all' prefix caching, " "please use '--mamba-cache-mode=align' instead" ) self.quant_config = vllm_config.quant_config super().__init__() self.config = config self.model = Qwen3NextMultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp") ) self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(config.vocab_size) self.set_moe_parameters() def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) def forward( self, input_ids: torch.Tensor | None, positions: torch.Tensor, hidden_states: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs: object, ): hidden_states = self.model( input_ids, positions, hidden_states, intermediate_tensors, inputs_embeds ) return hidden_states def compute_logits( self, hidden_states: torch.Tensor, spec_step_idx: int = 0, ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: shared_weight_names = ["embed_tokens", "lm_head"] def remap_weight_names(weights): for name, weight in weights: if name.startswith("mtp."): name = name.replace("mtp.", "model.") elif not any(key in name for key in shared_weight_names): continue yield name, weight loader = AutoWeightsLoader(self) return loader.load_weights(remap_weight_names(weights))