diff --git a/upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5.py b/upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5.py new file mode 100644 index 00000000..43b90046 --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5.py @@ -0,0 +1,819 @@ +# 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 Qwen3.5 Series compatible with HuggingFace weights.""" + +import typing +from collections.abc import Callable, Iterable + +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import ( + GemmaRMSNorm as Qwen3_5RMSNorm, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +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 vllm.transformers_utils.configs.qwen3_5 import ( + Qwen3_5Config, + Qwen3_5TextConfig, +) +from vllm.transformers_utils.configs.qwen3_5_moe import ( + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, +) + +from .interfaces import ( + HasInnerState, + IsHybrid, + MixtureOfExperts, + MultiModalEmbeddings, + SupportsEagle3, + SupportsLoRA, + SupportsPP, + _require_is_multimodal, +) +from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP +from .qwen3_next import ( + Qwen3NextAttention, + Qwen3NextDecoderLayer, + Qwen3NextModel, + Qwen3NextSparseMoeBlock, + QwenNextMixtureOfExperts, +) +from .qwen3_vl import ( + Qwen3_VisionTransformer, + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) +from .utils import ( + AutoWeightsLoader, + PPMissingLayer, + _merge_multimodal_embeddings, + extract_layer_index, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) + +logger = init_logger(__name__) + + +class Qwen3_5ProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(Qwen3_5Config) + + +class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(Qwen3_5MoeConfig) + + +class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): + def __init__( + self, + vllm_config: VllmConfig, + layer_type: str, + prefix: str = "", + ) -> None: + super(Qwen3NextDecoderLayer, self).__init__() + + config = vllm_config.model_config.hf_text_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.layer_type = layer_type + self.layer_idx = extract_layer_index(prefix) + + if self.layer_type == "linear_attention": + self.linear_attn = QwenGatedDeltaNetAttention( + config=config, + vllm_config=vllm_config, + prefix=f"{prefix}.linear_attn", + gqa_interleaved_layout=False, + ) + elif self.layer_type == "full_attention": + self.self_attn = Qwen3NextAttention( + config, + model_config=model_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + else: + raise ValueError(f"Invalid layer_type {self.layer_type}") + + # NOTE: Determine the MLP type based on the model type + # Qwen3.5 use all layers for MLP / Qwen3.5-MoE use sparse MoE blocks + if config.model_type == "qwen3_5_moe_text": + self.mlp = Qwen3NextSparseMoeBlock( + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + ) + elif config.model_type == "qwen3_5_text": + self.mlp = Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + raise ValueError(f"Invalid model_type {config.model_type}") + + self.input_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.layer_scale = getattr(config, "layer_scale", False) + if self.layer_scale: + self.attn_layer_scale = torch.nn.Parameter( + torch.zeros( + 1, + 1, + config.hidden_size, + ), + ) + self.ffn_layer_scale = torch.nn.Parameter( + torch.zeros( + 1, + 1, + config.hidden_size, + ), + ) + + +@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 Qwen3_5Model(Qwen3NextModel): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super(Qwen3NextModel, self).__init__() + + config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = ( + vllm_config.model_config.hf_text_config + ) + parallel_config = vllm_config.parallel_config + + eplb_config = parallel_config.eplb_config + self.num_redundant_experts = eplb_config.num_redundant_experts + + self.config = config + + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + ) + + def get_layer(prefix: str): + return Qwen3_5DecoderLayer( + vllm_config, + layer_type=config.layer_types[extract_layer_index(prefix)], + prefix=prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + if get_pp_group().is_last_rank: + self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.aux_hidden_state_layers: tuple[int, ...] = () + + 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=shard_id, + expert_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) + # GDN + ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)), + ("in_proj_qkvz", "in_proj_z", 3), + # self attention + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + # mlp + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("in_proj_ba", "in_proj_b", 0), + ("in_proj_ba", "in_proj_a", 1), + ] + + 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: list[tuple[str, str, int, str]] = [] + for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_up_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="gate_up_proj", + num_experts=1, + ): + if shard_id == "w3": + continue + parts = ckpt_name.split(".") + fused_expert_params_mapping.append( + (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id) + ) + num_experts = ( + self.config.num_experts if hasattr(self.config, "num_experts") else 0 + ) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if name.startswith("mtp."): + continue + + # Remapping the name of FP8 kv-scale. + if name.endswith("scale"): + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + 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 + + if weight_name not in name: + continue + + if "mlp.experts" in name: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip layers on other devices. + if is_pp_missing_parameter(name, self): + continue + # name = apply_attn_prefix(name, params_dict) + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + 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 + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + # Skip layers on other devices. + if is_pp_missing_parameter(name_mapped, self): + continue + if is_fused_expert: + # qwen3.5 no need to transpose + # loaded_weight = loaded_weight.transpose(-1, -2) + 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, + ) + if success: + name = name_mapped + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name_mapped.endswith(".bias") + or name_mapped.endswith("_bias") + ) and name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + weight_loader = 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 bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + if name not in params_dict: + logger.warning_once( + f"Parameter {name} not found in params_dict, skip loading" + ) + continue + 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 Qwen3_5ForCausalLMBase( + nn.Module, + HasInnerState, + SupportsEagle3, + SupportsLoRA, + SupportsPP, +): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": ["gate_proj", "up_proj"], + # GDN fused projections. + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + config = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + + scheduler_config = vllm_config.scheduler_config + if cache_config.mamba_cache_mode == "all": + raise NotImplementedError( + "Qwen3.5 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.scheduler_config = scheduler_config + self.model = Qwen3_5Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.model.aux_hidden_state_layers = layers + + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + num_layers = len(self.model.layers) + return (2, num_layers // 2, num_layers - 3) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ): + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=["mtp."], + ) + return loader.load_weights(weights) + + +class Qwen3_5ForCausalLM(Qwen3_5ForCausalLMBase): + pass + + +class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + + # set MoE hyperparameters + self.set_moe_parameters() + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() + + +######################################################## +# Qwen3_5-Dense +######################################################## + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid): + # Qwen3.5 does not support multimodal pruning (EVS). + supports_multimodal_pruning = False + + packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | { + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): + # protocols have not __init__ method, so we need to use nn.Module.__init__ + nn.Module.__init__(self) + config: Qwen3_5Config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + # Qwen3.5 does not support multimodal pruning (EVS). + self.is_multimodal_pruning_enabled = False + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = Qwen3_5ForCausalLM( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + is_multimodal = _require_is_multimodal(is_multimodal) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds + + def recompute_mrope_positions(self, *args, **kwargs): + raise NotImplementedError( + "Qwen3.5 does not support multimodal pruning (EVS). " + "recompute_mrope_positions should never be called." + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + """Run forward pass for Qwen3.5. + + Args: + input_ids: Flattened (concatenated) input_ids corresponding to a + batch. + positions: Flattened (concatenated) position ids corresponding to a + batch. + **NOTE**: If mrope is enabled (default setting for Qwen3VL + opensource models), the shape will be `(3, seq_len)`, + otherwise it will be `(seq_len,). + intermediate_tensors: Intermediate tensors from previous pipeline + stages. + inputs_embeds: Pre-computed input embeddings. + **kwargs: Additional keyword arguments including: + - pixel_values: Pixel values to be fed to a model. + `None` if no images are passed. + - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in + LLM. `None` if no images are passed. + - pixel_values_videos: Pixel values of videos to be fed to a + model. `None` if no videos are passed. + - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in + LLM. `None` if no videos are passed. + """ + + if intermediate_tensors is not None: + inputs_embeds = None + + hidden_states = self.language_model.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=["mtp."], + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + + +######################################################## +# Qwen3_5-MoE +######################################################## + + +class Qwen3_5_MoeMixtureOfExperts(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, Qwen3NextSparseMoeBlock): + 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 isinstance(layer, Qwen3_5DecoderLayer) and isinstance( + layer.mlp, Qwen3NextSparseMoeBlock + ): + example_moe = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_moe is None: + raise RuntimeError( + "No Qwen3_5 layer found in the language_model.model.layers." + ) + + # 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=Qwen3_5MoeProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class Qwen3_5MoeForConditionalGeneration( + Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts +): + # For MoE LoRA weights loading + is_3d_moe_weight: bool = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): + # protocols have not __init__ method, so we need to use nn.Module.__init__ + nn.Module.__init__(self) + config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + # Qwen3.5 does not support multimodal pruning (EVS). + self.is_multimodal_pruning_enabled = False + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = Qwen3_5MoeForCausalLM( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + # set MoE hyperparameters + self.set_moe_parameters() diff --git a/upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5_mtp.py b/upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5_mtp.py new file mode 100644 index 00000000..021462f3 --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5_mtp.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen3_5 MTP model.""" + +import typing +from collections.abc import Callable, Iterable + +import torch +from torch import nn + +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.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +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.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import LocalArgmaxMixin +from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5RMSNorm +from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig +from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig + +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + _require_is_multimodal, +) +from .utils import ( + AutoWeightsLoader, + PPMissingLayer, + _merge_multimodal_embeddings, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + maybe_prefix, +) + +logger = init_logger(__name__) + + +@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, + "hidden_states": 0, + } +) +class Qwen3_5MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + model_config = vllm_config.model_config + quant_config = vllm_config.quant_config + + config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = model_config.hf_text_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, "mtp_num_hidden_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 hf_quant_config.json exclude_modules. Force unquantized. + # Ref: https://github.com/vllm-project/vllm/pull/38650 + # Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124 + 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( + Qwen3_5DecoderLayer( + 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 = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.pre_fc_norm_hidden = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_embedding = Qwen3_5RMSNorm( + 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_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=shard_id, + expert_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), + ] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = fused_moe_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 + if hasattr(self.config, "num_experts") + else 0, + ) + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + is_fused_expert = False + fused_expert_params_mapping: list[tuple[str, str, int, str]] = [] + for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_up_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="gate_up_proj", + num_experts=1, + ): + if shard_id == "w3": + continue + parts = ckpt_name.split(".") + fused_expert_params_mapping.append( + (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id) + ) + num_experts = ( + self.config.num_experts if hasattr(self.config, "num_experts") else 0 + ) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + 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 + + if weight_name not in name: + continue + + if "mlp.experts" in name: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip layers on other devices. + if is_pp_missing_parameter(name, self): + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + 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 + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + # Skip layers on other devices. + if is_pp_missing_parameter(name_mapped, self): + continue + if is_fused_expert: + # qwen3.5 no need to transpose + # loaded_weight = loaded_weight.transpose(-1, -2) + 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, + ) + if success: + name = name_mapped + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name_mapped.endswith(".bias") + or name_mapped.endswith("_bias") + ) and name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + weight_loader = 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 bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + if name not in params_dict: + logger.warning_once( + f"Parameter {name} not found in params_dict, skip loading" + ) + continue + 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 + + +@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, + "hidden_states": 0, + } +) +class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal): + 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_text_config + self.vllm_config = vllm_config + cache_config = vllm_config.cache_config + if cache_config.mamba_cache_mode == "all": + raise NotImplementedError( + "Qwen3_5MTP 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 = Qwen3_5MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp") + ) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.model.embed_input_ids, + is_multimodal=is_multimodal, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + is_multimodal = _require_is_multimodal(is_multimodal) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds + + 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, + **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]: + def remap_weight_names(weights): + for name, weight in weights: + if name.startswith("mtp."): + name = name.replace("mtp.", "model.") + elif any(key in name for key in ["embed_tokens", "lm_head"]): + if "embed_tokens" in name: + name = name.replace("language_model.", "") + else: + continue + yield name, weight + + loader = AutoWeightsLoader(self) + return loader.load_weights(remap_weight_names(weights)) + + +class Qwen3_5MoeMTP(Qwen3_5MTP, QwenNextMixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + self.set_moe_parameters() diff --git a/upstream_ref/ds_vllm_latest/vllm/model_executor/models/registry.py b/upstream_ref/ds_vllm_latest/vllm/model_executor/models/registry.py new file mode 100644 index 00000000..722ba93d --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/model_executor/models/registry.py @@ -0,0 +1,1426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Whenever you add an architecture to this page, please also update +`tests/models/registry.py` with example HuggingFace models for it. +""" + +import importlib +import importlib.util +import json +import os +import pickle +import subprocess +import sys +import tempfile +from abc import ABC, abstractmethod +from collections.abc import Callable, Set +from dataclasses import asdict, dataclass, field +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING, Any, TypeVar + +import torch.nn as nn +import transformers + +from vllm import envs +from vllm.config import ( + ModelConfig, + iter_architecture_defaults, + try_match_architecture_defaults, +) +from vllm.logger import init_logger +from vllm.logging_utils import logtime +from vllm.tasks import ScoreType +from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module +from vllm.utils.hashing import safe_hash + +if TYPE_CHECKING: + from vllm.config.model import AttnTypeStr + from vllm.config.pooler import SequencePoolingType, TokenPoolingType +else: + AttnTypeStr = Any + SequencePoolingType = Any + TokenPoolingType = Any + + +from .interfaces import ( + has_inner_state, + has_noops, + is_attention_free, + is_hybrid, + requires_raw_input_tokens, + supports_mamba_prefix_caching, + supports_multimodal, + supports_multimodal_encoder_tp_data, + supports_multimodal_raw_input_only, + supports_pp, + supports_transcription, +) +from .interfaces_base import ( + get_attn_type, + get_default_seq_pooling_type, + get_default_tok_pooling_type, + get_score_type, + is_pooling_model, + is_text_generation_model, +) + +logger = init_logger(__name__) + +_TEXT_GENERATION_MODELS = { + # [Decoder-only] + "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"), + "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"), + "AquilaModel": ("llama", "LlamaForCausalLM"), + "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2 + "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"), + "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"), + "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"), + # baichuan-7b, upper case 'C' in the class name + "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), + # baichuan-13b, lower case 'c' in the class name + "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), + "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"), + "BailingMoeV2ForCausalLM": ("bailing_moe", "BailingMoeV2ForCausalLM"), + "BailingMoeV2_5ForCausalLM": ("bailing_moe_linear", "BailingMoeV25ForCausalLM"), + "BambaForCausalLM": ("bamba", "BambaForCausalLM"), + "BloomForCausalLM": ("bloom", "BloomForCausalLM"), + "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"), + "ChatGLMForConditionalGeneration": ("chatglm", "ChatGLMForCausalLM"), + "CohereForCausalLM": ("commandr", "CohereForCausalLM"), + "Cohere2ForCausalLM": ("commandr", "CohereForCausalLM"), + "Cohere2MoeForCausalLM": ("cohere2_moe", "Cohere2MoeForCausalLM"), + "CwmForCausalLM": ("llama", "LlamaForCausalLM"), + "DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"), + "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), + "DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"), + "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"), + "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"), + "DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"), + "DeepseekV4ForCausalLM": ("vllm.models.deepseek_v4", "DeepseekV4ForCausalLM"), + "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"), + "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"), + "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"), + "ExaoneForCausalLM": ("exaone", "ExaoneForCausalLM"), + "Exaone4ForCausalLM": ("exaone4", "Exaone4ForCausalLM"), + "ExaoneMoEForCausalLM": ("exaone_moe", "ExaoneMoeForCausalLM"), + "Fairseq2LlamaForCausalLM": ("fairseq2_llama", "Fairseq2LlamaForCausalLM"), + "FalconForCausalLM": ("falcon", "FalconForCausalLM"), + "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), + "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"), + "FlexOlmoForCausalLM": ("flex_olmo", "FlexOlmoForCausalLM"), + "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"), + "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"), + "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"), + "Rnj1ForCausalLM": ("rnj1", "Rnj1ForCausalLM"), + "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"), + "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"), + "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"), + "GlmForCausalLM": ("glm", "GlmForCausalLM"), + "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), + "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), + "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"), + "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"), + "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"), + "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"), + "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"), + "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"), + "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"), + "GraniteForCausalLM": ("granite", "GraniteForCausalLM"), + "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"), + "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), + "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), + "GritLM": ("gritlm", "GritLM"), + "Grok1ModelForCausalLM": ("grok1", "GrokForCausalLM"), + "Grok1ForCausalLM": ("grok1", "GrokForCausalLM"), + "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"), + "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"), + "HYV3ForCausalLM": ("hy_v3", "HYV3ForCausalLM"), + "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"), + "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"), + "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"), + "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"), + "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"), + "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"), + "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"), + "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), + "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"), + "JambaForCausalLM": ("jamba", "JambaForCausalLM"), + "KimiLinearForCausalLM": ("kimi_linear", "KimiLinearForCausalLM"), + "Lfm2ForCausalLM": ("lfm2", "Lfm2ForCausalLM"), + "Lfm2MoeForCausalLM": ("lfm2_moe", "Lfm2MoeForCausalLM"), + "LagunaForCausalLM": ("laguna", "LagunaForCausalLM"), + "LlamaForCausalLM": ("llama", "LlamaForCausalLM"), + "Llama4ForCausalLM": ("llama4", "Llama4ForCausalLM"), + # For decapoda-research/llama-* + "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"), + "LongcatFlashForCausalLM": ("longcat_flash", "LongcatFlashForCausalLM"), + "MambaForCausalLM": ("mamba", "MambaForCausalLM"), + "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"), + "MellumForCausalLM": ("mellum", "MellumForCausalLM"), + "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"), + "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"), + "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), + "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), + "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), + "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"), + "Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"), + "MistralForCausalLM": ("mistral", "MistralForCausalLM"), + "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"), + "MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"), + # transformers's mpt class has lower case + "MptForCausalLM": ("mpt", "MPTForCausalLM"), + "MPTForCausalLM": ("mpt", "MPTForCausalLM"), + "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"), + "MiMoV2FlashForCausalLM": ("mimo_v2", "MiMoV2FlashForCausalLM"), + "MiMoV2ForCausalLM": ("mimo_v2", "MiMoV2ForCausalLM"), + "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"), + "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), + "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), + "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"), + "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"), + "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"), + "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"), + "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), + "OPTForCausalLM": ("opt", "OPTForCausalLM"), + "OrionForCausalLM": ("orion", "OrionForCausalLM"), + "OuroForCausalLM": ("ouro", "OuroForCausalLM"), + "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"), + "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"), + "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"), + "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"), + "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"), + "PhiForCausalLM": ("phi", "PhiForCausalLM"), + "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), + "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"), + "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"), + "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"), + "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), + "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), + "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"), + "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"), + "RWForCausalLM": ("falcon", "FalconForCausalLM"), + "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), + "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), + "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"), + "Step1ForCausalLM": ("step1", "Step1ForCausalLM"), + "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"), + "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"), + "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"), + "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"), + "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"), + "SolarForCausalLM": ("solar", "SolarForCausalLM"), + "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat3ForCausalLM": ("llama", "LlamaForCausalLM"), + "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"), + "XverseForCausalLM": ("llama", "LlamaForCausalLM"), + "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"), +} + +_EMBEDDING_MODELS = { + # [Text-only] + "BertModel": ("bert", "BertEmbeddingModel"), + "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"), + "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"), + "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"), + "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"), + "Gemma3TextModel": ("gemma3", "Gemma3Model"), + "GlmForCausalLM": ("glm", "GlmForCausalLM"), + "GritLM": ("gritlm", "GritLM"), + "GteModel": ("bert_with_rope", "SnowflakeGteNewModel"), + "GteNewModel": ("bert_with_rope", "GteNewModel"), + "JinaEmbeddingsV5Model": ("jina", "JinaEmbeddingsV5Model"), + "LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"), + "LlamaModel": ("llama", "LlamaForCausalLM"), + **{ + # Multiple models share the same architecture, so we include them all + k: (mod, arch) + for k, (mod, arch) in _TEXT_GENERATION_MODELS.items() + if arch == "LlamaForCausalLM" + }, + "MistralModel": ("llama", "LlamaForCausalLM"), + "ModernBertModel": ("modernbert", "ModernBertModel"), + "NomicBertModel": ("bert_with_rope", "NomicBertModel"), + "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), + "Qwen2Model": ("qwen2", "Qwen2ForCausalLM"), + "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"), + "RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"), + "RobertaModel": ("roberta", "RobertaEmbeddingModel"), + "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "VoyageQwen3BidirectionalEmbedModel": ( + "voyage", + "VoyageQwen3BidirectionalEmbedModel", + ), + "XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"), + # [Multimodal] + "CLIPModel": ("clip", "CLIPEmbeddingModel"), + "ColPaliForRetrieval": ("colpali", "ColPaliModel"), + "LlamaNemotronVLModel": ("nemotron_vl", "LlamaNemotronVLForEmbedding"), + "LlavaNextForConditionalGeneration": ( + "llava_next", + "LlavaNextForConditionalGeneration", + ), + "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"), + "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), + "SiglipModel": ("siglip", "SiglipEmbeddingModel"), + # Technically Terratorch models work on images, both in + # input and output. I am adding it here because it piggy-backs on embedding + # models for the time being. + "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"), + "Terratorch": ("terratorch", "Terratorch"), +} + +_LATE_INTERACTION_MODELS = { + # [Text-only] + "HF_ColBERT": ("colbert", "ColBERTModel"), + "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"), + "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"), + "ColBERTLfm2Model": ("colbert", "ColBERTLfm2Model"), + "JinaForRanking": ("jina", "JinaForRanking"), + # [Multimodal] + "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"), + "ColPaliForRetrieval": ("colpali", "ColPaliModel"), + "ColQwen3": ("colqwen3", "ColQwen3Model"), + "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"), + "ColQwen3_5": ("colqwen3_5", "ColQwen3_5Model"), + "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"), +} + +_REWARD_MODELS = { + "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"), + "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"), + "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"), +} + +_TOKEN_CLASSIFICATION_MODELS = { + "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ModernBertForTokenClassification": ( + "modernbert", + "ModernBertForTokenClassification", + ), + "Qwen3ASRForcedAlignerForTokenClassification": ( + "qwen3_asr_forced_aligner", + "Qwen3ASRForcedAlignerForTokenClassification", + ), +} + +_SEQUENCE_CLASSIFICATION_MODELS = { + "BertForSequenceClassification": ("bert", "BertForSequenceClassification"), + "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"), + "GteNewForSequenceClassification": ( + "bert_with_rope", + "GteNewForSequenceClassification", + ), + "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), + "LlamaBidirectionalForSequenceClassification": ( + "llama", + "LlamaBidirectionalForSequenceClassification", + ), + "ModernBertForSequenceClassification": ( + "modernbert", + "ModernBertForSequenceClassification", + ), + "RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"), + "XLMRobertaForSequenceClassification": ( + "roberta", + "RobertaForSequenceClassification", + ), + # [Multimodal] + "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"), + "LlamaNemotronVLForSequenceClassification": ( + "nemotron_vl", + "LlamaNemotronVLForSequenceClassification", + ), +} + +_MULTIMODAL_MODELS = { + # [Decoder-only] + "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"), + "AudioFlamingo3ForConditionalGeneration": ( + "audioflamingo3", + "AudioFlamingo3ForConditionalGeneration", + ), + "MusicFlamingoForConditionalGeneration": ( + "musicflamingo", + "MusicFlamingoForConditionalGeneration", + ), + "AyaVisionForConditionalGeneration": ( + "aya_vision", + "AyaVisionForConditionalGeneration", + ), + "BagelForConditionalGeneration": ("bagel", "BagelForConditionalGeneration"), + "BeeForConditionalGeneration": ("bee", "BeeForConditionalGeneration"), + "Blip2ForConditionalGeneration": ("blip2", "Blip2ForConditionalGeneration"), + "ChameleonForConditionalGeneration": ( + "chameleon", + "ChameleonForConditionalGeneration", + ), + "Cheers": ("cheers", "CheersForConditionalGeneration"), + "CheersForConditionalGeneration": ("cheers", "CheersForConditionalGeneration"), + "Cohere2VisionForConditionalGeneration": ( + "cohere2_vision", + "Cohere2VisionForConditionalGeneration", + ), + "Cosmos3ForConditionalGeneration": ("cosmos3", "Cosmos3ForConditionalGeneration"), + "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"), + "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"), + "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"), + "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"), + "Eagle2_5_VLForConditionalGeneration": ( + "eagle2_5_vl", + "Eagle2_5_VLForConditionalGeneration", + ), + "Ernie4_5_VLMoeForConditionalGeneration": ( + "ernie45_vl", + "Ernie4_5_VLMoeForConditionalGeneration", + ), + "Exaone4_5_ForConditionalGeneration": ( + "exaone4_5", + "Exaone4_5_ForConditionalGeneration", + ), # noqa: E501 + "FireRedASR2ForConditionalGeneration": ( + "fireredasr2", + "FireRedASR2ForConditionalGeneration", + ), + "FunASRForConditionalGeneration": ("funasr", "FunASRForConditionalGeneration"), + "FireRedLIDForConditionalGeneration": ( + "fireredlid", + "FireRedLIDForConditionalGeneration", + ), + "FunAudioChatForConditionalGeneration": ( + "funaudiochat", + "FunAudioChatForConditionalGeneration", + ), + "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"), + "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), + "Gemma3nForConditionalGeneration": ( + "gemma3n_mm", + "Gemma3nForConditionalGeneration", + ), + "DiffusionGemmaForBlockDiffusion": ( + "diffusion_gemma", + "DiffusionGemmaForConditionalGeneration", + ), + "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"), + "Gemma4UnifiedForConditionalGeneration": ( + "gemma4_unified", + "Gemma4UnifiedForConditionalGeneration", + ), + "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"), + "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"), + "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), + "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"), + "GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"), + "GraniteSpeechForConditionalGeneration": ( + "granite_speech", + "GraniteSpeechForConditionalGeneration", + ), + "GraniteSpeechPlusForConditionalGeneration": ( + "granite_speech_plus", + "GraniteSpeechPlusForConditionalGeneration", + ), + "Granite4VisionForConditionalGeneration": ( + "granite4_vision", + "Granite4VisionForConditionalGeneration", + ), + "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"), + "HunYuanVLForConditionalGeneration": ( + "hunyuan_vision", + "HunYuanVLForConditionalGeneration", + ), + "InternVLChatModel": ("internvl", "InternVLChatModel"), + "InternS1ForConditionalGeneration": ( + "interns1", + "InternS1ForConditionalGeneration", + ), + "InternVLForConditionalGeneration": ( + "interns1", + "InternS1ForConditionalGeneration", + ), + "InternS1ProForConditionalGeneration": ( + "interns1_pro", + "InternS1ProForConditionalGeneration", + ), + "InternS2PreviewForConditionalGeneration": ( + "interns2_preview", + "InternS2PreviewForConditionalGeneration", + ), + "Idefics3ForConditionalGeneration": ( + "idefics3", + "Idefics3ForConditionalGeneration", + ), + "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"), + "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"), + "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"), + "KeyeVL1_5ForConditionalGeneration": ( + "keye_vl1_5", + "KeyeVL1_5ForConditionalGeneration", + ), + "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), + "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), + "MoonshotKimiaForCausalLM": ("kimi_audio", "KimiAudioForConditionalGeneration"), + "LightOnOCRForConditionalGeneration": ( + "lightonocr", + "LightOnOCRForConditionalGeneration", + ), + "Lfm2VlForConditionalGeneration": ("lfm2_vl", "Lfm2VLForConditionalGeneration"), + "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"), + "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"), + "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"), + "LlavaNextForConditionalGeneration": ( + "llava_next", + "LlavaNextForConditionalGeneration", + ), + "LlavaNextVideoForConditionalGeneration": ( + "llava_next_video", + "LlavaNextVideoForConditionalGeneration", + ), + "LlavaOnevisionForConditionalGeneration": ( + "llava_onevision", + "LlavaOnevisionForConditionalGeneration", + ), + "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"), + "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"), + "MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"), + "MiniMaxVL01ForConditionalGeneration": ( + "minimax_vl_01", + "MiniMaxVL01ForConditionalGeneration", + ), + "MiniCPMO": ("minicpmo", "MiniCPMO"), + "MiniCPMV": ("minicpmv", "MiniCPMV"), + "MiniCPMV4_6ForConditionalGeneration": ( + "minicpmv4_6", + "MiniCPMV4_6ForConditionalGeneration", + ), + "Mistral3ForConditionalGeneration": ( + "mistral3", + "Mistral3ForConditionalGeneration", + ), + "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"), + "Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"), + "Moondream3ForCausalLM": ("moondream3", "Moondream3ForCausalLM"), + "HfMoondream": ("moondream3", "Moondream3ForCausalLM"), + "NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "NVLM_D": ("nvlm_d", "NVLM_D_Model"), + "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"), + "OpenPanguVLForConditionalGeneration": ( + "openpangu_vl", + "OpenPanguVLForConditionalGeneration", + ), + "OpenVLAForActionPrediction": ("openvla", "OpenVLAForActionPrediction"), + "Ovis": ("ovis", "Ovis"), + "Ovis2_5": ("ovis2_5", "Ovis2_5"), + "Ovis2_6ForCausalLM": ("ovis2_5", "Ovis2_5"), + "Ovis2_6_MoeForCausalLM": ("ovis2_5", "Ovis2_5"), + "PaddleOCRVLForConditionalGeneration": ( + "paddleocr_vl", + "PaddleOCRVLForConditionalGeneration", + ), + "PaliGemmaForConditionalGeneration": ( + "paligemma", + "PaliGemmaForConditionalGeneration", + ), + "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"), + "Phi4ForCausalLMV": ("phi4siglip", "Phi4ForCausalLMV"), + "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"), + "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"), + "QianfanOCRForConditionalGeneration": ( + "qianfan_ocr", + "QianfanOCRForConditionalGeneration", + ), + "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), + "Qwen2_5_VLForConditionalGeneration": ( + "qwen2_5_vl", + "Qwen2_5_VLForConditionalGeneration", + ), + "Qwen2AudioForConditionalGeneration": ( + "qwen2_audio", + "Qwen2AudioForConditionalGeneration", + ), + "Qwen2_5OmniModel": ( + "qwen2_5_omni_thinker", + "Qwen2_5OmniThinkerForConditionalGeneration", + ), + "Qwen2_5OmniForConditionalGeneration": ( + "qwen2_5_omni_thinker", + "Qwen2_5OmniThinkerForConditionalGeneration", + ), + "Qwen3OmniMoeForConditionalGeneration": ( + "qwen3_omni_moe_thinker", + "Qwen3OmniMoeThinkerForConditionalGeneration", + ), + "Qwen3ASRForConditionalGeneration": ( + "qwen3_asr", + "Qwen3ASRForConditionalGeneration", + ), + "Qwen3ASRRealtimeGeneration": ("qwen3_asr_realtime", "Qwen3ASRRealtimeGeneration"), + "Qwen3VLForConditionalGeneration": ("qwen3_vl", "Qwen3VLForConditionalGeneration"), + "Qwen3VLMoeForConditionalGeneration": ( + "qwen3_vl_moe", + "Qwen3VLMoeForConditionalGeneration", + ), + "Qwen3_5ForConditionalGeneration": ("qwen3_5", "Qwen3_5ForConditionalGeneration"), + "Qwen3_5MoeForConditionalGeneration": ( + "qwen3_5", + "Qwen3_5MoeForConditionalGeneration", + ), + "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"), + "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"), + "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), + "StepVLForConditionalGeneration": ("step_vl", "StepVLForConditionalGeneration"), + "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"), + "Step3p7ForConditionalGeneration": ("step3p7", "Step3p7ForConditionalGeneration"), + "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), + "Tarsier2ForConditionalGeneration": ( + "qwen2_vl", + "Tarsier2ForConditionalGeneration", + ), + "UltravoxModel": ("ultravox", "UltravoxModel"), + "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), + "VoxtralRealtimeGeneration": ("voxtral_realtime", "VoxtralRealtimeGeneration"), + # [Encoder-decoder] + "CohereAsrForConditionalGeneration": ( + "cohere_asr", + "CohereAsrForConditionalGeneration", + ), + "NemotronParseForConditionalGeneration": ( + "nemotron_parse", + "NemotronParseForConditionalGeneration", + ), + "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"), +} + +_SPECULATIVE_DECODING_MODELS = { + "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"), + "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"), + "MiMoV2MTPModel": ("mimo_v2_mtp", "MiMoV2MTP"), + "MiMoV2OmniMTPModel": ("mimo_v2_mtp", "MiMoV2OmniMTP"), + "EagleCohereForCausalLM": ("cohere_eagle", "EagleCohereForCausalLM"), + "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"), + "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), + "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), + "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "EagleMistralForCausalLM": ("mistral_eagle", "EagleMistralForCausalLM"), + "EagleMistralLarge3ForCausalLM": ( + "mistral_large_3_eagle", + "EagleMistralLarge3ForCausalLM", + ), + "Eagle3DeepseekV2ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"), + "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"), + "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"), + "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"), + "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), + "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"), + "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"), + "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), + "Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"), + "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"), + "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"), + "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"), + "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"), + "GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"), + "MedusaModel": ("medusa", "Medusa"), + "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"), + "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"), + "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"), + "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), + "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), + "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"), + # Temporarily disabled. + # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1. + # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"), +} + +_TRANSFORMERS_SUPPORTED_MODELS = { + # Text generation models + "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"), + # Multimodal models + "Emu3ForConditionalGeneration": ( + "transformers", + "TransformersMultiModalForCausalLM", + ), +} + +_TRANSFORMERS_BACKEND_MODELS = { + # Text generation models + "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"), + "TransformersMoEForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + # Multimodal models + "TransformersMultiModalForCausalLM": ( + "transformers", + "TransformersMultiModalForCausalLM", + ), + "TransformersMultiModalMoEForCausalLM": ( + "transformers", + "TransformersMultiModalMoEForCausalLM", + ), + # Embedding models + "TransformersEmbeddingModel": ("transformers", "TransformersEmbeddingModel"), + "TransformersMoEEmbeddingModel": ("transformers", "TransformersMoEEmbeddingModel"), + "TransformersMultiModalEmbeddingModel": ( + "transformers", + "TransformersMultiModalEmbeddingModel", + ), + # Sequence classification models + "TransformersForSequenceClassification": ( + "transformers", + "TransformersForSequenceClassification", + ), + "TransformersMoEForSequenceClassification": ( + "transformers", + "TransformersMoEForSequenceClassification", + ), + "TransformersMultiModalForSequenceClassification": ( + "transformers", + "TransformersMultiModalForSequenceClassification", + ), +} + +_VLLM_MODELS = { + **_TEXT_GENERATION_MODELS, + **_EMBEDDING_MODELS, + **_LATE_INTERACTION_MODELS, + **_REWARD_MODELS, + **_TOKEN_CLASSIFICATION_MODELS, + **_SEQUENCE_CLASSIFICATION_MODELS, + **_MULTIMODAL_MODELS, + **_SPECULATIVE_DECODING_MODELS, + **_TRANSFORMERS_SUPPORTED_MODELS, + **_TRANSFORMERS_BACKEND_MODELS, +} + +# This variable is used as the args for subprocess.run(). We +# can modify this variable to alter the args if needed. e.g. +# when we use par format to pack things together, sys.executable +# might not be the target we want to run. +_SUBPROCESS_COMMAND = [sys.executable, "-m", "vllm.model_executor.models.registry"] + +_PREVIOUSLY_SUPPORTED_MODELS = { + "MotifForCausalLM": "0.10.2", + "Phi3SmallForCausalLM": "0.9.2", + "Phi4FlashForCausalLM": "0.10.2", + "Phi4MultimodalForCausalLM": "0.12.0", + "JAISLMHeadModel": "0.22.0", + "ErnieModel": "0.23.0", + "ErnieForSequenceClassification": "0.23.0", + "ErnieForTokenClassification": "0.23.0", + "QWenLMHeadModel": "0.23.0", + "QwenVLForConditionalGeneration": "0.23.0", + "InternLMForCausalLM": "0.23.0", + # encoder-decoder models except whisper + # have been removed for V0 deprecation. + "DonutForConditionalGeneration": "0.10.2", + "MllamaForConditionalGeneration": "0.10.2", +} + +_OOT_SUPPORTED_MODELS = { + "BartModel": "https://github.com/vllm-project/bart-plugin", + "BartForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", + "Florence2ForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", + "MBartForConditionalGeneration": "https://github.com/vllm-project/bart-plugin", +} + + +@dataclass(frozen=True) +class _ModelInfo: + architecture: str + is_text_generation_model: bool + is_pooling_model: bool + attn_type: AttnTypeStr + default_seq_pooling_type: SequencePoolingType + default_tok_pooling_type: TokenPoolingType + score_type: ScoreType + supports_multimodal: bool + supports_multimodal_raw_input_only: bool + requires_raw_input_tokens: bool + supports_multimodal_encoder_tp_data: bool + supports_pp: bool + has_inner_state: bool + is_attention_free: bool + is_hybrid: bool + has_noops: bool + supports_mamba_prefix_caching: bool + supports_transcription: bool + supports_transcription_only: bool + + @staticmethod + def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": + return _ModelInfo( + architecture=model.__name__, + is_text_generation_model=is_text_generation_model(model), + is_pooling_model=is_pooling_model(model), + default_seq_pooling_type=get_default_seq_pooling_type(model), + default_tok_pooling_type=get_default_tok_pooling_type(model), + attn_type=get_attn_type(model), + score_type=get_score_type(model), + supports_multimodal=supports_multimodal(model), + supports_multimodal_raw_input_only=supports_multimodal_raw_input_only( + model + ), + requires_raw_input_tokens=requires_raw_input_tokens(model), + supports_multimodal_encoder_tp_data=supports_multimodal_encoder_tp_data( + model + ), + supports_pp=supports_pp(model), + has_inner_state=has_inner_state(model), + is_attention_free=is_attention_free(model), + is_hybrid=is_hybrid(model), + supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), + supports_transcription=supports_transcription(model), + supports_transcription_only=( + supports_transcription(model) and model.supports_transcription_only + ), + has_noops=has_noops(model), + ) + + +class _BaseRegisteredModel(ABC): + @abstractmethod + def inspect_model_cls(self) -> _ModelInfo: + raise NotImplementedError + + @abstractmethod + def load_model_cls(self) -> type[nn.Module]: + raise NotImplementedError + + +@dataclass(frozen=True) +class _RegisteredModel(_BaseRegisteredModel): + """ + Represents a model that has already been imported in the main process. + """ + + interfaces: _ModelInfo + model_cls: type[nn.Module] + + @staticmethod + def from_model_cls(model_cls: type[nn.Module]): + return _RegisteredModel( + interfaces=_ModelInfo.from_model_cls(model_cls), + model_cls=model_cls, + ) + + def inspect_model_cls(self) -> _ModelInfo: + return self.interfaces + + def load_model_cls(self) -> type[nn.Module]: + return self.model_cls + + +@dataclass(frozen=True) +class _LazyRegisteredModel(_BaseRegisteredModel): + """ + Represents a model that has not been imported in the main process. + """ + + module_name: str + class_name: str + + @staticmethod + def _get_cache_dir() -> Path: + return Path(envs.VLLM_CACHE_ROOT) / "modelinfos" + + def _get_cache_filename(self) -> str: + cls_name = f"{self.module_name}-{self.class_name}".replace(".", "-") + return f"{cls_name}.json" + + def _load_modelinfo_from_cache(self, module_hash: str) -> _ModelInfo | None: + try: + try: + modelinfo_path = self._get_cache_dir() / self._get_cache_filename() + with open(modelinfo_path, encoding="utf-8") as file: + mi_dict = json.load(file) + except FileNotFoundError: + logger.debug( + "Cached model info file for class %s.%s not found", + self.module_name, + self.class_name, + ) + return None + + if mi_dict["hash"] != module_hash: + logger.debug( + "Cached model info file for class %s.%s is stale", + self.module_name, + self.class_name, + ) + return None + + # file not changed, use cached _ModelInfo properties + return _ModelInfo(**mi_dict["modelinfo"]) + except Exception: + logger.debug( + "Cached model info for class %s.%s error. ", + self.module_name, + self.class_name, + ) + return None + + def _save_modelinfo_to_cache(self, mi: _ModelInfo, module_hash: str) -> None: + """save dictionary json file to cache""" + from vllm.model_executor.model_loader.weight_utils import atomic_writer + + try: + modelinfo_dict = { + "hash": module_hash, + "modelinfo": asdict(mi), + } + cache_dir = self._get_cache_dir() + cache_dir.mkdir(parents=True, exist_ok=True) + modelinfo_path = cache_dir / self._get_cache_filename() + with atomic_writer(modelinfo_path, encoding="utf-8") as f: + json.dump(modelinfo_dict, f, indent=2) + except Exception: + logger.exception("Error saving model info cache.") + + @logtime(logger=logger, msg="Registry inspect model class") + def inspect_model_cls(self) -> _ModelInfo: + # Modules registered with a non-default location (e.g. the + # hardware-isolated ``vllm.models.`` layout) live outside + # ``vllm/model_executor/models``. Resolve the module spec directly + # so the file-hash cache stays warm for them. + if self.module_name.startswith("vllm.model_executor.models."): + model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py" + else: + try: + spec = importlib.util.find_spec(self.module_name) + except (ImportError, ValueError): + spec = None + model_path = Path(spec.origin) if spec is not None and spec.origin else None + module_hash = None + + if model_path is not None and model_path.exists(): + with open(model_path, "rb") as f: + module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest() + + mi = self._load_modelinfo_from_cache(module_hash) + if mi is not None: + logger.debug( + "Loaded model info for class %s.%s from cache", + self.module_name, + self.class_name, + ) + return mi + else: + logger.debug( + "Cache model info for class %s.%s miss. Loading model instead.", + self.module_name, + self.class_name, + ) + + # Performed in another process to avoid initializing CUDA + mi = _run_in_subprocess( + lambda: _ModelInfo.from_model_cls(self.load_model_cls()) + ) + logger.debug( + "Loaded model info for class %s.%s", self.module_name, self.class_name + ) + + # save cache file + if module_hash is not None: + self._save_modelinfo_to_cache(mi, module_hash) + + return mi + + def load_model_cls(self) -> type[nn.Module]: + mod = importlib.import_module(self.module_name) + return getattr(mod, self.class_name) + + +@lru_cache(maxsize=128) +def _try_load_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> type[nn.Module] | None: + from vllm.platforms import current_platform + + current_platform.verify_model_arch(model_arch) + try: + return model.load_model_cls() + except Exception: + logger.exception("Error in loading model architecture '%s'", model_arch) + return None + + +@lru_cache(maxsize=128) +def _try_inspect_model_cls( + model_arch: str, + model: _BaseRegisteredModel, +) -> _ModelInfo | None: + try: + return model.inspect_model_cls() + except Exception: + logger.exception("Error in inspecting model architecture '%s'", model_arch) + return None + + +@dataclass +class _ModelRegistry: + # Keyed by model_arch + models: dict[str, _BaseRegisteredModel] = field(default_factory=dict) + + def get_supported_archs(self) -> Set[str]: + return self.models.keys() + + def register_model( + self, + model_arch: str, + model_cls: type[nn.Module] | str, + ) -> None: + """ + Register an external model to be used in vLLM. + + `model_cls` can be either: + + - A [`torch.nn.Module`][] class directly referencing the model. + - A string in the format `:` which can be used to + lazily import the model. This is useful to avoid initializing CUDA + when importing the model and thus the related error + `RuntimeError: Cannot re-initialize CUDA in forked subprocess`. + """ + if not isinstance(model_arch, str): + msg = f"`model_arch` should be a string, not a {type(model_arch)}" + raise TypeError(msg) + + if model_arch in self.models: + logger.debug( + "Model architecture %s is already registered, and will be " + "overwritten by the new model class %s.", + model_arch, + model_cls, + ) + + if isinstance(model_cls, str): + split_str = model_cls.split(":") + if len(split_str) != 2: + msg = "Expected a string in the format `:`" + raise ValueError(msg) + + model = _LazyRegisteredModel(*split_str) + elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module): + model = _RegisteredModel.from_model_cls(model_cls) + else: + msg = ( + "`model_cls` should be a string or PyTorch model class, " + f"not a {type(model_arch)}" + ) + raise TypeError(msg) + + self.models[model_arch] = model + + def _raise_for_unsupported(self, architectures: list[str]): + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + for arch in architectures: + if arch in _PREVIOUSLY_SUPPORTED_MODELS: + previous_version = _PREVIOUSLY_SUPPORTED_MODELS[arch] + + raise ValueError( + f"Model architecture {arch} was supported in vLLM until " + f"v{previous_version}, and is not supported anymore. " + "Please use an older version of vLLM if you want to " + "use this model architecture." + ) + if arch in _OOT_SUPPORTED_MODELS: + plugin_url = _OOT_SUPPORTED_MODELS[arch] + + raise ValueError( + f"Model architecture {arch} is not supported in-tree anymore. " + f"Please install the plugin at {plugin_url} if you want to " + "use this model architecture." + ) + + raise ValueError( + f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: + if model_arch not in self.models: + return None + + return _try_load_model_cls(model_arch, self.models[model_arch]) + + def _try_inspect_model_cls(self, model_arch: str) -> _ModelInfo | None: + if model_arch not in self.models: + return None + + return _try_inspect_model_cls(model_arch, self.models[model_arch]) + + def _try_resolve_transformers( + self, + architecture: str, + model_config: ModelConfig, + ) -> str | None: + if architecture in _TRANSFORMERS_BACKEND_MODELS: + return architecture + + auto_map: dict[str, str] = ( + getattr(model_config.hf_config, "auto_map", None) or dict() + ) + + # Make sure that config class is always initialized before model class, + # otherwise the model class won't be able to access the config class, + # the expected auto_map should have correct order like: + # "auto_map": { + # "AutoConfig": "--", + # "AutoModel": "--", + # "AutoModelFor": "--", + # }, + for prefix in ("AutoConfig", "AutoModel"): + for name, module in auto_map.items(): + if name.startswith(prefix): + try_get_class_from_dynamic_module( + module, + model_config.model, + revision=model_config.revision, + code_revision=model_config.code_revision, + trust_remote_code=model_config.trust_remote_code, + warn_on_fail=False, + ) + + model_module = getattr(transformers, architecture, None) + + if model_module is None: + for name, module in auto_map.items(): + if name.startswith("AutoModel"): + model_module = try_get_class_from_dynamic_module( + module, + model_config.model, + revision=model_config.revision, + code_revision=model_config.code_revision, + trust_remote_code=model_config.trust_remote_code, + warn_on_fail=True, + ) + if model_module is not None: + break + else: + if model_config.model_impl != "transformers": + return None + + raise ValueError( + f"Cannot find model module. {architecture!r} is not a " + "registered model in the Transformers library (only " + "relevant if the model is meant to be in Transformers) " + "and 'AutoModel' is not present in the model config's " + "'auto_map' (relevant if the model is custom)." + ) + + if not model_module.is_backend_compatible(): + if model_config.model_impl != "transformers": + return None + + raise ValueError( + f"The Transformers implementation of {architecture!r} " + "is not compatible with vLLM." + ) + + return model_config._get_transformers_backend_cls() + + def _normalize_arch( + self, + architecture: str, + model_config: ModelConfig, + ) -> str: + if architecture in self.models: + return architecture + + # This may be called in order to resolve runner_type and convert_type + # in the first place, in which case we consider the default match + match = try_match_architecture_defaults( + architecture, + runner_type=getattr(model_config, "runner_type", None), + convert_type=getattr(model_config, "convert_type", None), + ) + if match: + suffix, _ = match + + # Get the name of the base model to convert + for repl_suffix, _ in iter_architecture_defaults(): + base_arch = architecture.replace(suffix, repl_suffix) + if base_arch in self.models: + return base_arch + + return architecture + + def inspect_model_cls( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> tuple[_ModelInfo, str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + raise ValueError("No model architectures are specified") + + # Require transformers impl + if model_config.model_impl == "transformers": + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + elif model_config.model_impl == "terratorch": + model_info = self._try_inspect_model_cls("Terratorch") + return (model_info, "Terratorch") + + # Fallback to transformers impl (after resolving convert_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + and getattr(model_config, "convert_type", "none") == "none" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + + for arch in architectures: + normalized_arch = self._normalize_arch(arch, model_config) + model_info = self._try_inspect_model_cls(normalized_arch) + if model_info is not None: + return (model_info, arch) + + # Fallback to transformers impl (before resolving runner_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_info = self._try_inspect_model_cls(arch) + if model_info is not None: + return (model_info, arch) + + return self._raise_for_unsupported(architectures) + + def resolve_model_cls( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> tuple[type[nn.Module], str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + raise ValueError("No model architectures are specified") + + # Require transformers impl + if model_config.model_impl == "transformers": + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + elif model_config.model_impl == "terratorch": + arch = "Terratorch" + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + # Fallback to transformers impl (after resolving convert_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + and getattr(model_config, "convert_type", "none") == "none" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + for arch in architectures: + normalized_arch = self._normalize_arch(arch, model_config) + model_cls = self._try_load_model_cls(normalized_arch) + if model_cls is not None: + return (model_cls, arch) + + # Fallback to transformers impl (before resolving runner_type) + if ( + all(arch not in self.models for arch in architectures) + and model_config.model_impl == "auto" + ): + arch = self._try_resolve_transformers(architectures[0], model_config) + if arch is not None: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + return self._raise_for_unsupported(architectures) + + def is_text_generation_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_text_generation_model + + def is_pooling_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_pooling_model + + def is_multimodal_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_multimodal + + def is_multimodal_raw_input_only_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_multimodal_raw_input_only + + def is_pp_supported_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_pp + + def model_has_inner_state( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.has_inner_state + + def is_attention_free_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_attention_free + + def is_hybrid_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.is_hybrid + + def is_noops_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.has_noops + + def is_transcription_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_transcription + + def is_transcription_only_model( + self, + architectures: str | list[str], + model_config: ModelConfig, + ) -> bool: + model_cls, _ = self.inspect_model_cls(architectures, model_config) + return model_cls.supports_transcription_only + + +def _resolve_module_name(mod_relname: str) -> str: + # Allow registry entries to point at fully-qualified module paths (e.g. + # ``vllm.models.deepseek_v4``) for models that live outside the legacy + # ``vllm.model_executor.models`` flat layout. + if mod_relname.startswith("vllm."): + return mod_relname + return f"vllm.model_executor.models.{mod_relname}" + + +ModelRegistry = _ModelRegistry( + { + model_arch: _LazyRegisteredModel( + module_name=_resolve_module_name(mod_relname), + class_name=cls_name, + ) + for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items() + } +) + +_T = TypeVar("_T") + + +def _run_in_subprocess(fn: Callable[[], _T]) -> _T: + # NOTE: We use a temporary directory instead of a temporary file to avoid + # issues like https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file + with tempfile.TemporaryDirectory() as tempdir: + output_filepath = os.path.join(tempdir, "registry_output.tmp") + + # `cloudpickle` allows pickling lambda functions directly + import cloudpickle + + input_bytes = cloudpickle.dumps((fn, output_filepath)) + + # cannot use `sys.executable __file__` here because the script + # contains relative imports + returned = subprocess.run( + _SUBPROCESS_COMMAND, input=input_bytes, capture_output=True + ) + + # check if the subprocess is successful + try: + returned.check_returncode() + except Exception as e: + # wrap raised exception to provide more information + raise RuntimeError( + f"Error raised in subprocess:\n{returned.stderr.decode()}" + ) from e + + with open(output_filepath, "rb") as f: + return pickle.load(f) + + +def _run() -> None: + # Setup plugins + from vllm.plugins import load_general_plugins + + load_general_plugins() + + fn, output_file = pickle.loads(sys.stdin.buffer.read()) + + result = fn() + + with open(output_file, "wb") as f: + f.write(pickle.dumps(result)) + + +if __name__ == "__main__": + _run() diff --git a/upstream_ref/ds_vllm_latest/vllm/multimodal/__init__.py b/upstream_ref/ds_vllm_latest/vllm/multimodal/__init__.py new file mode 100644 index 00000000..34438a4f --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/multimodal/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .hasher import MultiModalHasher +from .inputs import BatchedTensorInputs, MultiModalKwargsItems, NestedTensors +from .registry import MultiModalRegistry + +MULTIMODAL_REGISTRY = MultiModalRegistry() +""" +The global [`MultiModalRegistry`][vllm.multimodal.registry.MultiModalRegistry] +is used by model runners to dispatch data processing according to the target +model. + +Info: + [mm_processing](../../../design/mm_processing.md) +""" + +__all__ = [ + "BatchedTensorInputs", + "MultiModalHasher", + "MultiModalKwargsItems", + "NestedTensors", + "MULTIMODAL_REGISTRY", + "MultiModalRegistry", +] diff --git a/upstream_ref/ds_vllm_latest/vllm/multimodal/registry.py b/upstream_ref/ds_vllm_latest/vllm/multimodal/registry.py new file mode 100644 index 00000000..6fdae470 --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/multimodal/registry.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading +from collections import defaultdict +from collections.abc import Mapping +from dataclasses import dataclass +from multiprocessing.synchronize import Lock as LockType +from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar, cast + +from vllm.inputs import MultiModalInput +from vllm.logger import init_logger +from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config + +from .cache import ( + BaseMultiModalProcessorCache, + BaseMultiModalReceiverCache, + MultiModalProcessorOnlyCache, + MultiModalProcessorSenderCache, + MultiModalReceiverCache, + ShmObjectStoreReceiverCache, + ShmObjectStoreSenderCache, +) +from .processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + InputProcessingContext, + TimingContext, +) + +if TYPE_CHECKING: + from vllm.config import ModelConfig, ObservabilityConfig, VllmConfig + from vllm.model_executor.models.interfaces import SupportsMultiModal + +logger = init_logger(__name__) + +N = TypeVar("N", bound=type["SupportsMultiModal"]) +_I = TypeVar("_I", bound=BaseProcessingInfo) +_I_co = TypeVar("_I_co", bound=BaseProcessingInfo, covariant=True) + + +class ProcessingInfoFactory(Protocol[_I_co]): + """ + Constructs a + [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor] + instance from the context. + """ + + def __call__( + self, + ctx: InputProcessingContext, + ) -> _I_co: ... + + +class DummyInputsBuilderFactory(Protocol[_I]): # type: ignore[misc] + """ + Constructs a + [`BaseDummyInputsBuilder`][vllm.multimodal.processing.BaseDummyInputsBuilder] + instance from the context. + """ + + def __call__(self, info: _I) -> BaseDummyInputsBuilder[_I]: ... + + +class MultiModalProcessorFactory(Protocol[_I]): # type: ignore[misc] + """ + Constructs a + [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor] + instance from the context. + """ + + def __call__( + self, + info: _I, + dummy_inputs: BaseDummyInputsBuilder[_I], + *, + cache: BaseMultiModalProcessorCache | None = None, + ) -> BaseMultiModalProcessor[_I]: ... + + +@dataclass(frozen=True) +class _ProcessorFactories(Generic[_I]): + info: ProcessingInfoFactory[_I] + processor: MultiModalProcessorFactory[_I] + dummy_inputs: DummyInputsBuilderFactory[_I] + + def build_processor( + self, + ctx: InputProcessingContext, + *, + cache: BaseMultiModalProcessorCache | None = None, + ): + info = self.info(ctx) + dummy_inputs_builder = self.dummy_inputs(info) + return self.processor(info, dummy_inputs_builder, cache=cache) + + +class MultiModalRegistry: + """ + A registry that dispatches data processing according to the model. + """ + + def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool: + """ + Checks if the model supports multimodal inputs. + Returns True if the model is multimodal with any non-zero supported + modalities, otherwise returns False, effectively running in + text-only mode. + """ + if not model_config.is_multimodal_model: + return False + + mm_config = model_config.get_multimodal_config() + try: + info = self._create_processing_info(model_config, tokenizer=None) + except ValueError: + logger.warning_once( + "Model %s is treated as multimodal but has no registered " + "multimodal processor; running in text-only mode.", + model_config.model, + ) + return False + + # Check if all supported modalities have limit == 0 + if all( + mm_config.get_limit_per_prompt(modality) == 0 + for modality in info.supported_mm_limits + ): + # If enable_mm_embeds is True, we still need MM infrastructure + # to process pre-computed embeddings even though encoder won't run + if mm_config.enable_mm_embeds: + return True + + logger.info_once( + "All limits of multimodal modalities supported by the model " + "are set to 0, running in text-only mode." + ) + return False + + return True + + def register_processor( + self, + processor: MultiModalProcessorFactory[_I], + *, + info: ProcessingInfoFactory[_I], + dummy_inputs: DummyInputsBuilderFactory[_I], + ): + """ + Register a multi-modal processor to a model class. The processor + is constructed lazily, hence a factory method should be passed. + + When the model receives multi-modal data, the provided function is + invoked to transform the data into a dictionary of model inputs. + """ + + def wrapper(model_cls: N) -> N: + if "_processor_factory" in model_cls.__dict__: + logger.warning( + "Model class %s already has a multi-modal processor " + "registered to %s. It is overwritten by the new one.", + model_cls, + self, + ) + + model_cls._processor_factory = _ProcessorFactories( + info=info, + dummy_inputs=dummy_inputs, + processor=processor, + ) + + return model_cls + + return wrapper + + def _get_model_cls(self, model_config: "ModelConfig") -> "SupportsMultiModal": + # Avoid circular import + from vllm.model_executor.model_loader import get_model_architecture + + model_cls, _ = get_model_architecture(model_config) + if not hasattr(model_cls, "_processor_factory"): + raise ValueError( + f"Model class {model_cls.__name__} has no registered " + "multimodal processor" + ) + return cast("SupportsMultiModal", model_cls) + + def _create_processing_ctx( + self, + model_config: "ModelConfig", + tokenizer: TokenizerLike | None = None, + ) -> InputProcessingContext: + if tokenizer is None: + tokenizer = cached_tokenizer_from_config(model_config) + + return InputProcessingContext(model_config, tokenizer) + + def _create_processing_info( + self, + model_config: "ModelConfig", + tokenizer: TokenizerLike | None = None, + ) -> BaseProcessingInfo: + model_cls = self._get_model_cls(model_config) + factories = model_cls._processor_factory + ctx = self._create_processing_ctx(model_config, tokenizer) + return factories.info(ctx) + + def get_processing_info(self, model_config: "ModelConfig") -> BaseProcessingInfo: + return self._create_processing_info(model_config, tokenizer=None) + + def create_processor( + self, + model_config: "ModelConfig", + *, + tokenizer: TokenizerLike | None = None, + cache: BaseMultiModalProcessorCache | None = None, + ) -> BaseMultiModalProcessor[BaseProcessingInfo]: + """ + Create a multi-modal processor for a specific model and tokenizer. + """ + if not model_config.is_multimodal_model: + model_name = model_config.served_model_name or model_config.model + raise ValueError(f"{model_name} is not a multimodal model") + + model_cls = self._get_model_cls(model_config) + factories = model_cls._processor_factory + + ctx = self._create_processing_ctx(model_config, tokenizer) + + return factories.build_processor(ctx, cache=cache) + + def get_dummy_mm_inputs( + self, + model_config: "ModelConfig", + mm_counts: Mapping[str, int], + *, + cache: BaseMultiModalProcessorCache | None = None, + processor: BaseMultiModalProcessor | None = None, + ) -> MultiModalInput: + """ + Create dummy data for profiling the memory usage of a model. + + The model is identified by `model_config`. + """ + seq_len = model_config.max_model_len + + if processor is None: + processor = self.create_processor(model_config, cache=cache) + + mm_config = model_config.get_multimodal_config() + processor_inputs = processor.dummy_inputs.get_dummy_processor_inputs( + seq_len=seq_len, + mm_counts=mm_counts, + mm_options=mm_config.limit_per_prompt, + ) + mm_inputs = processor.apply( + processor_inputs, + timing_ctx=TimingContext(enabled=False), + ) + + prompt_token_ids = mm_inputs["prompt_token_ids"] + total_len = len(prompt_token_ids) + if total_len < seq_len: + prompt_token_ids.extend([0] * (seq_len - total_len)) + + return mm_inputs + + def _get_cache_type( + self, + vllm_config: "VllmConfig", + ) -> Literal[None, "processor_only", "lru", "shm"]: + model_config = vllm_config.model_config + if not self.supports_multimodal_inputs(model_config): + return None + + # Check if the cache is disabled. + mm_config = model_config.get_multimodal_config() + if mm_config.mm_processor_cache_gb <= 0: + return None + + # Check if IPC caching is supported. + parallel_config = vllm_config.parallel_config + is_ipc_supported = parallel_config._api_process_count == 1 and ( + parallel_config.data_parallel_size == 1 + or parallel_config.data_parallel_external_lb + ) + + if not is_ipc_supported: + return "processor_only" + + mm_config = model_config.get_multimodal_config() + return mm_config.mm_processor_cache_type + + def processor_cache_from_config( + self, + vllm_config: "VllmConfig", + ) -> BaseMultiModalProcessorCache | None: + """Return a `BaseMultiModalProcessorCache`, if enabled.""" + cache_type = self._get_cache_type(vllm_config) + if cache_type is None: + return None + elif cache_type == "processor_only": + return MultiModalProcessorOnlyCache(vllm_config.model_config) + elif cache_type == "lru": + return MultiModalProcessorSenderCache(vllm_config.model_config) + elif cache_type == "shm": + return ShmObjectStoreSenderCache(vllm_config) + else: + raise ValueError(f"Unknown cache type: {cache_type!r}") + + def processor_only_cache_from_config( + self, + vllm_config: "VllmConfig", + ) -> MultiModalProcessorOnlyCache | None: + """Return a `MultiModalProcessorOnlyCache`, if enabled.""" + cache_type = self._get_cache_type(vllm_config) + if cache_type is None: + return None + + return MultiModalProcessorOnlyCache(vllm_config.model_config) + + def engine_receiver_cache_from_config( + self, + vllm_config: "VllmConfig", + ) -> BaseMultiModalReceiverCache | None: + """Return a `BaseMultiModalReceiverCache` for the engine process.""" + cache_type = self._get_cache_type(vllm_config) + if cache_type in (None, "processor_only", "shm"): + return None + elif cache_type == "lru": + return MultiModalReceiverCache(vllm_config.model_config) + else: + raise ValueError(f"Unknown cache type: {cache_type!r}") + + def worker_receiver_cache_from_config( + self, + vllm_config: "VllmConfig", + shared_worker_lock: LockType, + ) -> BaseMultiModalReceiverCache | None: + """Return a `BaseMultiModalReceiverCache` for the worker process.""" + cache_type = self._get_cache_type(vllm_config) + if cache_type in (None, "processor_only", "lru"): + return None + elif cache_type == "shm": + return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock) + else: + raise ValueError(f"Unknown cache type: {cache_type!r}") + + +class MultiModalTimingRegistry: + def __init__(self, observability_config: "ObservabilityConfig | None") -> None: + super().__init__() + + if observability_config and observability_config.enable_mm_processor_stats: + self._lock = threading.Lock() + self._ctx_by_request_id = defaultdict[str, TimingContext](TimingContext) + self._enabled = True + else: + self._enabled = False + + def get(self, request_id: str) -> TimingContext: + if not self._enabled: + return TimingContext(enabled=False) + + with self._lock: + return self._ctx_by_request_id[request_id] + + def stat(self) -> dict[str, dict[str, float]]: + if not self._enabled: + return {} + + with self._lock: + stats = { + req_id: ctx.get_stats_dict() + for req_id, ctx in self._ctx_by_request_id.items() + } + self._ctx_by_request_id.clear() + return stats diff --git a/upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5.py b/upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5.py new file mode 100644 index 00000000..d5820a57 --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. +# All rights reserved. +# +# 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. +"""Qwen3.5 model configuration""" + +from transformers.configuration_utils import PretrainedConfig + + +class Qwen3_5TextConfig(PretrainedConfig): + model_type = "qwen3_5_text" + keys_to_ignore_at_inference = ["past_key_values"] + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + base_config_key = "text_config" + + def __init__( + self, + vocab_size=248320, + hidden_size=4096, + intermediate_size=12288, + num_hidden_layers=32, + num_attention_heads=16, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_parameters=None, + attention_bias=False, + attention_dropout=0.0, + head_dim=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + layer_types=None, + pad_token_id=None, + bos_token_id=None, + eos_token_id=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.head_dim = head_dim + self.rope_parameters = rope_parameters + kwargs.setdefault("partial_rotary_factor", 0.25) + + self.layer_types = layer_types + if self.layer_types is None: + interval_pattern = kwargs.get("full_attention_interval", 4) + self.layer_types = [ + "linear_attention" + if bool((i + 1) % interval_pattern) + else "full_attention" + for i in range(self.num_hidden_layers) + ] + kwargs["ignore_keys_at_rope_validation"] = { + "mrope_section", + "mrope_interleaved", + } + self.validate_layer_type() + + # linear attention part + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + super().__init__(**kwargs) + # Set these AFTER super().__init__() because transformers v4's + # PretrainedConfig.__init__ has these as explicit params with different + # defaults (e.g. tie_word_embeddings=True) that would overwrite our values. + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.tie_word_embeddings = tie_word_embeddings + + +class Qwen3_5VisionConfig(PretrainedConfig): + model_type = "qwen3_5" + base_config_key = "vision_config" + + def __init__( + self, + depth=27, + hidden_size=1152, + hidden_act="gelu_pytorch_tanh", + intermediate_size=4304, + num_heads=16, + in_channels=3, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, + out_hidden_size=3584, + num_position_embeddings=2304, + initializer_range=0.02, + **kwargs, + ): + super().__init__(**kwargs) + + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.num_position_embeddings = num_position_embeddings + self.initializer_range = initializer_range + + +class Qwen3_5Config(PretrainedConfig): + model_type = "qwen3_5" + sub_configs = { + "vision_config": Qwen3_5VisionConfig, + "text_config": Qwen3_5TextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + tie_word_embeddings=False, + **kwargs, + ): + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + + if isinstance(text_config, dict): + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + self.text_config = self.sub_configs["text_config"]() + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + super().__init__(**kwargs) + # Set after super().__init__() to avoid v4 PretrainedConfig overwrite + self.tie_word_embeddings = tie_word_embeddings + + +__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig"] diff --git a/upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5_moe.py b/upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5_moe.py new file mode 100644 index 00000000..ec229ce8 --- /dev/null +++ b/upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5_moe.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. +# All rights reserved. +# +# 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. +"""Qwen3.5-MoE model configuration""" + +from transformers.configuration_utils import PretrainedConfig + + +class Qwen3_5MoeTextConfig(PretrainedConfig): + model_type = "qwen3_5_moe_text" + keys_to_ignore_at_inference = ["past_key_values"] + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.shared_expert.gate_proj": "colwise", + "layers.*.mlp.shared_expert.up_proj": "colwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + base_config_key = "text_config" + + def __init__( + self, + vocab_size=248320, + hidden_size=2048, + num_hidden_layers=40, + num_attention_heads=16, + num_key_value_heads=2, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_parameters=None, + attention_bias=False, + attention_dropout=0.0, + head_dim=256, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + moe_intermediate_size=512, + shared_expert_intermediate_size=512, + num_experts_per_tok=8, + num_experts=256, + output_router_logits=False, + router_aux_loss_coef=0.001, + layer_types=None, + pad_token_id=None, + bos_token_id=None, + eos_token_id=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.head_dim = head_dim + self.rope_parameters = rope_parameters + kwargs.setdefault("partial_rotary_factor", 0.25) + + self.layer_types = layer_types + if self.layer_types is None: + interval_pattern = kwargs.get("full_attention_interval", 4) + self.layer_types = [ + "linear_attention" + if bool((i + 1) % interval_pattern) + else "full_attention" + for i in range(self.num_hidden_layers) + ] + kwargs["ignore_keys_at_rope_validation"] = { + "mrope_section", + "mrope_interleaved", + } + self.validate_layer_type() + + # linear attention part + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + super().__init__(**kwargs) + # Set these AFTER super().__init__() because transformers v4's + # PretrainedConfig.__init__ has these as explicit params with different + # defaults (e.g. tie_word_embeddings=True) that would overwrite our values. + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.tie_word_embeddings = tie_word_embeddings + + +class Qwen3_5MoeVisionConfig(PretrainedConfig): + model_type = "qwen3_5_moe" + base_config_key = "vision_config" + + def __init__( + self, + depth=27, + hidden_size=1152, + hidden_act="gelu_pytorch_tanh", + intermediate_size=4304, + num_heads=16, + in_channels=3, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, + out_hidden_size=3584, + num_position_embeddings=2304, + initializer_range=0.02, + **kwargs, + ): + super().__init__(**kwargs) + + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.num_position_embeddings = num_position_embeddings + self.initializer_range = initializer_range + + +class Qwen3_5MoeConfig(PretrainedConfig): + model_type = "qwen3_5_moe" + sub_configs = { + "vision_config": Qwen3_5MoeVisionConfig, + "text_config": Qwen3_5MoeTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + tie_word_embeddings=False, + **kwargs, + ): + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + + if isinstance(text_config, dict): + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + self.text_config = self.sub_configs["text_config"]() + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + super().__init__(**kwargs) + # Set after super().__init__() to avoid v4 PretrainedConfig overwrite + self.tie_word_embeddings = tie_word_embeddings + + +__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"] diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp new file mode 100644 index 00000000..1ad364a4 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp @@ -0,0 +1,32 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "ilu_ops_api.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode) { + if (act_mode == "silu") { + infer::silu_and_mul(input, out); + } else { + LOG(FATAL) << "Unsupported act mode: " << act_mode + << ", only support silu, gelu, gelu_tanh"; + } +} +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp new file mode 100644 index 00000000..ad3cd295 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp @@ -0,0 +1,163 @@ + +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "ixinfer.h" +#include "utils.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void reshape_paged_cache(torch::Tensor& key, + std::optional& value, + torch::Tensor& key_cache, + std::optional& value_cache, + torch::Tensor& slot_mapping) { + auto value_ = value.value_or(torch::Tensor()); + auto value_cache_ = value_cache.value_or(torch::Tensor()); + + int64_t key_token_stride = key.stride(0); + int64_t value_token_stride = 0; + if (value_.defined()) { + value_token_stride = value_.stride(0); + } + slot_mapping = slot_mapping.to(at::kLong); + infer::xllm_reshape_and_cache(key, + value_, + key_cache, + value_cache_, + slot_mapping, + key_token_stride, + value_token_stride); +} + +void batch_prefill(torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& v_quant_scale, + const torch::Tensor& block_tables, + int64_t max_query_len, + int64_t max_seq_len, + float scale, + bool is_causal, + int64_t window_size_left, + int64_t window_size_right, + const std::string& compute_dtype, + bool return_lse) { + double softcap = 0.0; + bool sqrt_alibi = false; + auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor()); + auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor()); + auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor()); + auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor()); + auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor()); + auto block_tables_ = block_tables; + auto key_ = key; + auto value_ = value.value(); + infer::ixinfer_flash_attn_unpad_with_block_tables(query, + key_, + value_, + output, + block_tables_, + q_cu_seq_lens_, + kv_cu_seq_lens_, + max_query_len, + max_seq_len, + is_causal, + window_size_left, + window_size_right, + static_cast(scale), + softcap, + sqrt_alibi, + alibi_slope, + c10::nullopt, + output_lse); +} + +void batch_decode(torch::Tensor& query, + const torch::Tensor& k_cache, + torch::Tensor& output, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const std::optional& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& mask, + const std::string& compute_dtype, + int64_t max_seq_len, + int64_t window_size_left, + int64_t window_size_right, + float scale, + bool return_lse, + bool is_causal, + int64_t kv_cache_quant_bit_size) { + if (query.dim() == 4) { + query = + query + .view({query.size(0) * query.size(1), query.size(2), query.size(3)}) + .contiguous(); + } + if (output.dim() == 4) { + output = output + .view({output.size(0) * output.size(1), + output.size(2), + output.size(3)}) + .contiguous(); + ; + } + auto v_cache_ = v_cache.value_or(torch::Tensor()); + int64_t num_kv_heads = k_cache.size(1); + int64_t page_block_size = k_cache.size(2); + double softcap = 0.0; + bool enable_cuda_graph = false; + bool use_sqrt_alibi = false; + auto block_table_ = block_table; + auto k_cache_ = k_cache; + auto seq_lens_ = seq_lens; + infer::xllm_paged_attention(output, + query, + k_cache_, + v_cache_, + num_kv_heads, + scale, + block_table_, + seq_lens_, + page_block_size, + max_seq_len, + alibi_slope, + is_causal, + (int32_t)window_size_left, + (int32_t)window_size_right, + softcap, + enable_cuda_graph, + use_sqrt_alibi, + c10::nullopt); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp new file mode 100644 index 00000000..21c15d8c --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp @@ -0,0 +1,99 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include + +#include "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias) { + torch::Tensor input_ = input.to(torch::kFloat32); + auto reduce_weight = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kFloat).device(input.device())); + auto topk_indices = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kInt32).device(input.device())); + auto token_expert_indices = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kInt32).device(input.device())); + + infer::topk_softmax( + reduce_weight, topk_indices, token_expert_indices, input_, false); + + auto tt = reduce_weight.sum(-1); + if (normalize) { + reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1); + } + return std::make_tuple(reduce_weight, topk_indices); +} + +std::vector moe_gen_idx(torch::Tensor& expert_id, + int64_t expert_num) { + auto src_dst = expert_id.new_empty({expert_id.numel()}); + auto dst_src = torch::empty_like(src_dst); + auto expert_sizes_gpu = expert_id.new_empty({expert_num}); + auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1}); + infer::moe_compute_token_index_api(expert_id, + src_dst, + dst_src, + expert_sizes_gpu, + /*expert_mask=*/std::nullopt, + /*expert_sizes_cpu*/ std::nullopt, + /*expert_sizes_gpu*/ std::nullopt, + 0, + expert_num, + expert_num); + + expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1); + return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum}; +} + +torch::Tensor moe_expand_input(const torch::Tensor& input, + const torch::Tensor& gather_index, + const torch::Tensor& combine_idx, + int64_t topk) { + int64_t dst_tokens = input.size(0) * topk; + auto output = input.new_empty({dst_tokens, input.size(1)}); + infer::moe_expand_input( + output, input, combine_idx, gather_index, dst_tokens, topk); + + return output; +} + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) { + input = input.view({-1, weight.size(1), input.size(1)}); + auto output = input.new_empty({input.size(0), input.size(2)}); + infer::moe_output_reduce_sum(output, + input, + weight, + /*mask=*/std::nullopt, + /*extra_residual*/ std::nullopt, + /*scaling_factor=*/1.0); + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp new file mode 100644 index 00000000..290299a0 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp @@ -0,0 +1,39 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +torch::Tensor group_gemm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& tokens_per_experts, + const std::optional& dst_to_src, + torch::Tensor& output) { + infer::moe_w16a16_group_gemm( + output, + input, + weight, + tokens_per_experts, + dst_to_src, + /*bias=*/std::nullopt, + /*format=*/"TN", + /*persistent=*/0, + /*output_n=*/tokens_per_experts.sum().item()); + + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h b/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h new file mode 100644 index 00000000..3dedd7da --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h @@ -0,0 +1,153 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "ATen/Tensor.h" +#include "ATen/cuda/CUDAEvent.h" +#include "c10/core/Device.h" +#include "c10/core/DeviceGuard.h" +#include "c10/core/GradMode.h" +#include "c10/core/InferenceMode.h" +#include "c10/core/MemoryFormat.h" +#include "c10/core/ScalarType.h" +#include "c10/core/TensorOptions.h" +#include "c10/cuda/CUDAFunctions.h" +#include "c10/cuda/CUDAGuard.h" +#include "c10/cuda/CUDAStream.h" +#include "ixformer.h" +#include "kernels/kernels.h" + +// #include "utils.h" +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& cos_sin_cache, + torch::Tensor& positions, + bool interleave); + +// act_mode only support silu, gelu, gelu_tanh +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode); + +void reshape_paged_cache( + torch::Tensor& key, // (num_tokens, num_heads, head_size) + std::optional& value, // (num_tokens, num_heads, head_size) + torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size) + std::optional& + value_cache, // (num_blocks, num_heads, block_size, head_size) + torch::Tensor& slot_mapping); //(num_tokens) + +void batch_prefill(torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& v_quant_scale, + const torch::Tensor& block_tables, + int64_t max_query_len, + int64_t max_seq_len, + float scale, + bool is_causal, + int64_t window_size_left, + int64_t window_size_right, + const std::string& compute_dtype, + bool return_lse); + +void batch_decode(torch::Tensor& query, + const torch::Tensor& k_cache, + torch::Tensor& output, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const std::optional& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& mask, + const std::string& compute_dtype, + int64_t max_seq_len, + int64_t window_size_left, + int64_t window_size_right, + float scale, + bool return_lse, + bool is_causal, + int64_t kv_cache_quant_bit_size); + +void residual_layer_norm(torch::Tensor& input, + torch::Tensor& output, + std::optional& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& residual_out, + double eps); + +void rms_norm(torch::Tensor& output, + torch::Tensor& input, + torch::Tensor& weight, + double eps); + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias); + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias); + +std::vector moe_gen_idx(torch::Tensor& expert_id, + int64_t expert_num); + +torch::Tensor moe_expand_input(const torch::Tensor& input, + const torch::Tensor& gather_index, + const torch::Tensor& combine_idx, + int64_t topk); + +torch::Tensor group_gemm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& tokens_per_experts, + const std::optional& dst_to_src, + torch::Tensor& output); + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight); +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h b/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h new file mode 100644 index 00000000..83bad88e --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h @@ -0,0 +1,147 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ +#include + +#include "ATen/Tensor.h" +#include "utils.h" + +namespace ixformer::infer { +torch::Tensor ixinfer_flash_attn_unpad_with_block_tables( + torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& out, + torch::Tensor& block_tables, + torch::Tensor& cu_seq_q, + torch::Tensor& cu_seq_k, + int64_t max_seq_q, + int64_t max_seq_k, + bool is_causal, + int64_t window_left, + int64_t window_right, + double scale, + double softcap, + bool sqrt_alibi, + const std::optional& alibi_slopes, + const std::optional& sinks, + std::optional& lse); + +void silu_and_mul(torch::Tensor& input, torch::Tensor& output); + +torch::Tensor xllm_paged_attention( + torch::Tensor& out, + torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + int64_t num_kv_heads, + double scale, + torch::Tensor& block_tables, + torch::Tensor& context_lens, + int64_t block_size, + int64_t max_context_len, + const std::optional& alibi_slopes, + bool causal, + int32_t window_left, + int32_t window_right, + double softcap, + bool enable_cuda_graph, + bool use_sqrt_alibi, + const std::optional& sinks); + +torch::Tensor ixformer_linear(torch::Tensor& input, + torch::Tensor& weight, + int64_t act_type, + const std::optional& bias, + const std::optional& out, + const std::optional persistent); + +torch::Tensor ixformer_linear_ex(torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& out); + +void xllm_reshape_and_cache(torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& slot_mapping, + int64_t key_token_stride, + int64_t value_token_stride); + +void xllm_rotary_embedding(torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int64_t head_size, + torch::Tensor& cos_sin_cache, + bool is_neox); + +void residual_rms_norm(torch::Tensor& input, + torch::Tensor& residual, + torch::Tensor& weight, + torch::Tensor& output, + torch::Tensor& residual_output, + const std::optional& fused_bias, + double alpha, + double eps, + bool is_post); + +void rms_norm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& output, + const std::optional& fused_bias, + double eps); + +void topk_softmax(torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& token_expert_indices, + torch::Tensor& gating_output, + bool renormalize); + +void moe_compute_token_index_api( + torch::Tensor& topk_ids, + torch::Tensor& src_dst, + torch::Tensor& dst_src, + torch::Tensor& expert_sizes_gpu, + const c10::optional& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& expand_tokens_gpu, + int64_t start_expert_id, + int64_t end_expert_id, + int64_t num_experts); + +void moe_expand_input(torch::Tensor outputs, + torch::Tensor inputs, + torch::Tensor dst_to_src, + const c10::optional& src_to_dst, + int64_t dst_tokens, + int64_t expand_factor); + +void moe_w16a16_group_gemm(torch::Tensor output, + torch::Tensor inputs, + torch::Tensor weights, + torch::Tensor tokens_per_experts, + const c10::optional& dst_to_src, + const c10::optional& bias, + std::string format, + int64_t persistent, + int64_t output_n); + +void moe_output_reduce_sum(torch::Tensor outputs, + torch::Tensor inputs, + const c10::optional& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, + double scaling_factor); +} // namespace ixformer::infer diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp new file mode 100644 index 00000000..f90c0d47 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp @@ -0,0 +1,73 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "util/env_var.h" + +namespace xllm::kernel::ilu { + +bool gemv_conditions(const torch::Tensor& input, + const torch::Tensor& weight, + const torch::Tensor& bias, + int64_t gemv_max_batch) { + // gemv input:[m,k] weight:[n,k] + // 1. m <= gemv_max_batch + // 2. k % 32 == 0 && n % 2 == 0 + // 3. bias is None + + torch::Tensor input_view = input.view({-1, input.size(-1)}); + torch::Tensor weight_view = weight.view({-1, weight.size(-1)}); + + int64_t m = input_view.size(0); + int64_t k = input_view.size(1); + int64_t n = weight_view.size(0); + + if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 && + n % 2 == 0) { + return true; + } + return false; +} + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias) { + int64_t act_type = -1; + bool persistent = false; + std::vector output_shape = a.sizes().vec(); + if (!output_shape.empty()) { + output_shape[output_shape.size() - 1] = b.size(0); + } + torch::Tensor output = a.new_empty(output_shape); + + bool use_gemv = true; + const int64_t gemv_max_batch = 1; + const bool disable_infer_gemm_ex = + xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false); + + use_gemv = + use_gemv && + gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) && + !disable_infer_gemm_ex && (act_type == -1); + + if (use_gemv) { + output = infer::ixformer_linear_ex(a, b, bias, output); + } else { + output = infer::ixformer_linear(a, b, act_type, bias, output, persistent); + } + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp new file mode 100644 index 00000000..e451bc36 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp @@ -0,0 +1,51 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "utils.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void residual_layer_norm(torch::Tensor& input, + torch::Tensor& output, + std::optional& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& residual_out, + double eps) { + auto residual_ = residual.value_or(torch::zeros_like(input)); + torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input)); + infer::residual_rms_norm(input, + residual_, + weight, + output, + residual_out_, + bias, + /*alpha=*/1.0, + eps, + false); +} + +void rms_norm(torch::Tensor& output, + torch::Tensor& input, + torch::Tensor& weight, + double eps) { + std::optional fused_bias = std::nullopt; + infer::rms_norm(input, weight, output, fused_bias, eps); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp b/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp new file mode 100644 index 00000000..45af7656 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp @@ -0,0 +1,31 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "ilu_ops_api.h" +#include "utils.h" + +namespace xllm::kernel::ilu { + +void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& cos_sin_cache, + torch::Tensor& positions, + bool interleave) { + const int64_t head_size = cos_sin_cache.size(-1); + infer::xllm_rotary_embedding( + positions, query, key, head_size, cos_sin_cache, !interleave); +} + +} // namespace xllm::kernel::ilu diff --git a/upstream_ref/xllm_latest/core/kernels/ilu/utils.h b/upstream_ref/xllm_latest/core/kernels/ilu/utils.h new file mode 100644 index 00000000..9fd15298 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/ilu/utils.h @@ -0,0 +1,63 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ +#pragma once +namespace xllm::kernel::ilu { +#undef check_tensor_contiguous +#define check_tensor_contiguous(x, type) \ + TORCH_CHECK(x.scalar_type() == type); \ + TORCH_CHECK(x.is_cuda()); \ + TORCH_CHECK(x.is_contiguous()); + +#undef check_tensor_half_bf_float +#define check_tensor_half_bf_float(x) \ + TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || \ + x.scalar_type() == at::ScalarType::Float || \ + x.scalar_type() == at::ScalarType::BFloat16); \ + TORCH_CHECK(x.is_cuda()); + +// from torchCheckMsgImpl +inline const char* ixformer_check_msg_impl(const char* msg) { return msg; } +// // If there is just 1 user-provided C-string argument, use it. + +#define IXFORMER_CHECK_MSG(cond, type, ...) \ + (ixformer_check_msg_impl( \ + "Expected " #cond \ + " to be true, but got false. " \ + "(Could this error message be improved? If so, " \ + "please report an enhancement request to ixformer.)", \ + ##__VA_ARGS__)) + +#define IXFORMER_CHECK(cond, ...) \ + { \ + if (!(cond)) { \ + std::cerr << __FILE__ << " (" << __LINE__ << ")" \ + << "-" << __FUNCTION__ << " : " \ + << IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \ + throw std::runtime_error("IXFORMER_CHECK ERROR"); \ + } \ + } + +#undef CUINFER_CHECK +#define CUINFER_CHECK(func) \ + do { \ + cuinferStatus_t status = (func); \ + if (status != CUINFER_STATUS_SUCCESS) { \ + std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \ + << ": " << cuinferGetErrorString(status) << std::endl; \ + throw std::runtime_error("CUINFER_CHECK ERROR"); \ + } \ + } while (0) + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp b/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp new file mode 100644 index 00000000..62d36a62 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/ilu/attention.cpp @@ -0,0 +1,189 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "attention.h" + +#include "kernels/ilu/ilu_ops_api.h" +#include "kernels/ops_api.h" + +namespace xllm { +namespace layer { +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window) + : num_heads_(num_heads), + head_size_(head_size), + scale_(scale), + num_kv_heads_(num_kv_heads), + v_head_dim_(head_size), + use_fused_mla_qkv_(false), + enable_lighting_indexer_(false), + enable_mla_(false), + sliding_window_(sliding_window) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + int64_t num_kv_heads, + int64_t v_head_dim, + int64_t sliding_window, + float scale, + bool use_fused_mla_qkv, + bool enable_lighting_indexer, + bool enable_mla) + : num_heads_(num_heads), + head_size_(head_size), + scale_(scale), + num_kv_heads_(num_kv_heads), + v_head_dim_(v_head_dim), + use_fused_mla_qkv_(use_fused_mla_qkv), + enable_lighting_indexer_(enable_lighting_indexer), + enable_mla_(enable_mla), + sliding_window_(sliding_window) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +std::tuple> AttentionImpl::forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache) { + std::optional output_lse = std::nullopt; + torch::Tensor output; + if (enable_mla_) { + output = torch::empty({query.size(0), num_heads_ * v_head_dim_}, + query.options()); + } else { + output = torch::empty_like(query); + } + if (attn_metadata.is_dummy) { + return std::make_tuple(output, output_lse); + } + + bool only_prefill = + attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_; + torch::Tensor k_cache = kv_cache.get_k_cache(); + std::optional v_cache; + std::optional v; + if (!enable_mla_) { + v = value.view({-1, num_kv_heads, head_size_}); + v_cache = kv_cache.get_v_cache(); + } + + bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_); + if (!skip_process_cache) { + xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; + reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_}); + reshape_paged_cache_params.value = v; + reshape_paged_cache_params.k_cache = k_cache; + reshape_paged_cache_params.v_cache = v_cache; + reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; + xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); + } + + if (enable_lighting_indexer_ || !only_prefill) { + decoder_forward(query, output, k_cache, v_cache, attn_metadata); + } else { + prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata); + } + + int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_; + output = output.view({-1, num_heads_ * head_size}); + return {output, output_lse}; +} + +void AttentionImpl::prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + std::optional output_lse = std::nullopt; + query = query.view({-1, num_heads_, head_size_}); + output = output.view({-1, num_heads_, head_size_v}); + // torch::Tensor k_cache_ = k_cache; + // torch::Tensor v_cache_ = v_cache.value(); + xllm::kernel::ilu::batch_prefill(query, + k_cache, + v_cache, + output, + output_lse, + attn_metadata.q_cu_seq_lens, + attn_metadata.kv_cu_seq_lens, + /*alibi_slope=*/std::nullopt, + /*attn_bias=*/std::nullopt, + /*q_quant_scale=*/std::nullopt, + /*k_quant_scale=*/std::nullopt, + /*v_quant_scale=*/std::nullopt, + attn_metadata.block_table, + attn_metadata.max_query_len, + attn_metadata.max_seq_len, + scale_, + attn_metadata.is_causal, + sliding_window_, + /*window_size_right=*/-1, + attn_metadata.compute_dtype, + /*return_lse=*/false); +} + +void AttentionImpl::decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + query = query.view({-1, 1, num_heads_, head_size_}); + output = output.view({-1, 1, num_heads_, head_size_v}); + std::optional output_lse = std::nullopt; + + int64_t block_aligned_max_seq_len = + attn_metadata.block_table.size(-1) * k_cache.size(2); + + xllm::kernel::ilu::batch_decode(query, + k_cache, + output, + attn_metadata.block_table, + attn_metadata.kv_seq_lens, + v_cache, + output_lse, + /*q_quant_scale=*/std::nullopt, + /*k_quant_scale=*/std::nullopt, + /*v_quant_scale=*/std::nullopt, + /*out_quant_scale=*/std::nullopt, + /*alibi_slope=*/std::nullopt, + attn_metadata.attn_mask, + attn_metadata.compute_dtype, + block_aligned_max_seq_len, + sliding_window_, + /*window_size_right=*/-1, + scale_, + /*return_lse=*/false, + attn_metadata.is_causal, + /*kv_cache_quant_bit_size=*/-1); +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/ilu/attention.h b/upstream_ref/xllm_latest/core/layers/ilu/attention.h new file mode 100644 index 00000000..bf4b59ba --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/ilu/attention.h @@ -0,0 +1,82 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_input_params.h" +#include "layers/common/attention_metadata.h" + +namespace xllm { +namespace layer { +class AttentionImpl : public torch::nn::Module { + public: + AttentionImpl() = default; + + AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window); + AttentionImpl(int64_t num_heads, + int64_t head_size, + int64_t num_kv_heads, + int64_t v_head_dim, + int64_t sliding_window, + float scale, + bool use_fused_mla_qkv, + bool enable_lighting_indexer, + bool enable_mla); + + std::tuple> forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache); + + void prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + void decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + private: + int64_t num_heads_; + int64_t head_size_; + float scale_; + int64_t num_kv_heads_; + int64_t v_head_dim_; + bool use_fused_mla_qkv_; + bool enable_lighting_indexer_; + bool enable_mla_; + int64_t sliding_window_; +}; +TORCH_MODULE(Attention); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp new file mode 100644 index 00000000..7c829ad3 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp @@ -0,0 +1,806 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "fused_moe.h" + +#include + +#include + +#include "common/global_flags.h" +#include "core/framework/config/eplb_config.h" +#include "core/framework/config/scheduler_config.h" +#include "core/framework/config/speculative_config.h" +#include "framework/parallel_state/parallel_state.h" +#include "kernels/ops_api.h" +#include "layers/common/dp_utils.h" +#include "util/utils.h" + +namespace { + +int32_t get_dtype_size(torch::ScalarType dtype) { + return static_cast(torch::elementSize(dtype)); +} + +} // namespace + +namespace xllm { +namespace layer { + +FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : num_total_experts_(static_cast(model_args.n_routed_experts())), + topk_(model_args.num_experts_per_tok()), + num_expert_group_(model_args.n_group()), + topk_group_(model_args.topk_group()), + route_scale_(model_args.routed_scaling_factor()), + hidden_size_(model_args.hidden_size()), + n_shared_experts_(model_args.n_shared_experts()), + is_gated_(moe_args.is_gated), + renormalize_(model_args.norm_topk_prob() ? 1 : 0), + hidden_act_(model_args.hidden_act()), + scoring_func_(model_args.scoring_func()), + quant_args_(quant_args), + parallel_args_(parallel_args), + options_(options), + device_(options.device()) { + const int64_t num_experts = num_total_experts_; + const int64_t intermediate_size = + static_cast(model_args.moe_intermediate_size()); + const std::string& topk_method = model_args.topk_method(); + int64_t ep_size = parallel_args.ep_size(); + int64_t ep_rank = 0; + tp_pg_ = parallel_args.tp_group_; + if (ep_size > 1) { + ep_rank = parallel_args.moe_ep_group_->rank(); + tp_pg_ = parallel_args.moe_tp_group_; + } + + // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is + // supported + if (!quant_args.quant_method().empty()) { + if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 || + !quant_args.activation_dynamic()) { + LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when " + "quant_method is set. " + << "Got quant_method=" << quant_args.quant_method() + << ", bits=" << quant_args.bits() + << ", activation_dynamic=" << quant_args.activation_dynamic(); + } + // If confirmed as smoothquant w8a8, set is_smoothquant_ to true + is_smoothquant_ = true; + } else { + is_smoothquant_ = false; + } + + // Deep EP initialization check + enable_deep_ep_ = + ::xllm::EPLBConfig::get_instance().expert_parallel_degree() == 2 && + ep_size > 1; + if (enable_deep_ep_) { + // for now, we only implement the deep ep for decode stage. + // so we will assume the max_token_num is limited to max_batch_size * (1+K) + // K is the number of speculative tokens. + int64_t dispatch_token_size; + if (quant_args.quant_method() == "smoothquant") { + // float32 is for the scale of the quantized input + dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) + + get_dtype_size(torch::kFloat32); + } else { + dispatch_token_size = + hidden_size_ * get_dtype_size(options_.dtype().toScalarType()); + } + torch::ScalarType combine_dtype = options_.dtype().toScalarType(); + int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype); + // Ensure calculation base is at least ep_size + int64_t effective_seqs = std::max( + (int64_t)::xllm::SchedulerConfig::get_instance().max_seqs_per_batch(), + (int64_t)ep_size); + // NOTE: ::xllm::SchedulerConfig::get_instance().max_seqs_per_batch() + // represents the maximum total batch size, regardless of the dp size. To + // ensure robust scheduling and account for the worst-case scenario, we must + // guarantee that each rank is capable of handling the maximum possible + // number of tokens. Therefore, we define max_num_tokens_per_rank as the + // full maximum value, without dividing by either the rank count or the dp + // size. + int64_t max_num_tokens_per_rank = + (1 + + ::xllm::SpeculativeConfig::get_instance().num_speculative_tokens()) * + effective_seqs * topk_; + + // make sure that all layers share the same deep ep instance + // so that the memory footprint is minimized + deep_ep_ = DeepEPManager::get_instance(dispatch_token_size, + combine_token_size, + max_num_tokens_per_rank, + num_experts, + parallel_args, + options_); + + // obtain the buffer and parameters of deep ep + deep_ep_buffer_ = deep_ep_->get_buffer(); + deep_ep_params_ = deep_ep_->get_params(); + + // intermediate buffer that can be initialized once + // we place these tensor here in order to speed up forward pass + int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv; + int64_t token_bytes = is_smoothquant_ + ? get_dtype_size(torch::kInt8) + : get_dtype_size(options_.dtype().toScalarType()); + token_bytes = token_bytes * hidden_size_; + int64_t head_size = n_tokens_recv * token_bytes; + dispatch_recv_token_tensor_head_ = + deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size) + .view({n_tokens_recv, token_bytes}); + // input scale in smoothquant + if (is_smoothquant_) { + int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32); + dispatch_recv_token_tensor_tail_ = + deep_ep_buffer_.combine_send_token_tensor + .narrow(0, head_size, tail_size) + .view({n_tokens_recv, -1}); + } + } + + // calculate the number of experts per rank + num_experts_per_rank_ = num_experts / ep_size; + start_expert_id_ = ep_rank * num_experts_per_rank_; + + if (topk_method == "noaux_tc") { + e_score_correction_bias_ = register_parameter( + "e_score_correction_bias", torch::empty({num_experts}, options), false); + } + + gate_ = register_module( + "gate_proj", + ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options)); + if (n_shared_experts_ > 0) { + ProcessGroup* shared_expert_pg; + if (parallel_args_.ep_size() > 1) { + // we use tp=1 for shared experts computation in deep ep mode + CHECK(parallel_args_.ep_size() == parallel_args_.world_size()) + << "Models with shared experts only support ep_size equal to " + "world size for now."; + shared_expert_pg = parallel_args.moe_tp_group_; + } else { + shared_expert_pg = parallel_args.process_group_; + } + // The shared experts computation can proceed in parallel with the + // final communication step during the MoE computation, as long as it + // remains independent of any communication operations. For optimal + // performance, ensure that the shared experts layer on each rank always + // maintains its own unique weights. + shared_experts_ = + register_module("shared_experts", + DenseMLP(hidden_size_, + intermediate_size * n_shared_experts_, + is_gated_, + false, + hidden_act_, + /*enable_result_reduction=*/true, + quant_args, + shared_expert_pg, + options)); + } + + // create weight buffer + const int64_t world_size = tp_pg_->world_size(); + int64_t local_intermediate_size = intermediate_size / world_size; + if (is_smoothquant_) { + auto quant_option = options_.dtype(torch::kInt8); + auto fp_option = options_.dtype(torch::kFloat32); + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + quant_option), + false); + w13_scale_ = register_parameter( + "w13_scale", + torch::empty({num_experts_per_rank_, local_intermediate_size * 2}, + fp_option), + false); + // Note: We do not check enable_deep_ep_ here, since smooth quantization + // information may be needed even when deep EP mode is disabled. This allows + // retrieving quantization parameters for any subset of experts as required. + input_smooth_ = register_parameter( + "input_smooth", + torch::empty({num_total_experts_, hidden_size_}, fp_option), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + quant_option), + false); + w2_scale_ = register_parameter( + "w2_scale", + torch::empty({num_experts_per_rank_, hidden_size_}, fp_option), + false); + act_smooth_ = register_parameter( + "act_smooth", + torch::empty({num_experts_per_rank_, local_intermediate_size}, + fp_option), + false); + + } else { + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + options_), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + options_), + false); + } +} + +torch::Tensor FusedMoEImpl::create_group_gemm_output( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype, + torch::Tensor& workspace) { + // unify shape logic: define the target shape once. + bool is_3d_weight = (b.dim() != 2); + int64_t num_tokens = a.size(0); + int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0); + + std::vector output_shape; + int64_t required_elements = num_tokens * out_dim; + + if (is_3d_weight) { + output_shape = {num_tokens, out_dim}; + } else { + output_shape = {group_list.size(0), num_tokens, out_dim}; + required_elements *= group_list.size(0); + } + + auto options = a.options().dtype(dtype); + + // non-smoothquant: direct allocation + if (!is_smoothquant_) { + return torch::empty(output_shape, options); + } + + // smoothquant: managed workspace logic + if (!workspace.defined()) { + // Lazy initialization: allocate max buffer for the lifecycle + // Note: accessing class members w13_ and w2_ directly for context + int64_t max_width = std::max(w13_.size(1), w2_.size(1)); + workspace = torch::empty({num_tokens * max_width}, options); + } + + // view construction + CHECK(workspace.numel() >= required_elements) + << "FusedMoE Workspace too small! Alloc: " << workspace.numel() + << ", Req: " << required_elements; + + // utilize the pre-calculated output_shape + return workspace.slice(0, 0, required_elements).view(output_shape); +} + +torch::Tensor FusedMoEImpl::select_experts( + const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info, + bool enable_all2all_communication) { + // prepare the parameters for select_experts + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + int64_t expert_size = w13_.size(0); + + // Step 1: apply softmax topk or sigmoid topk / routing logic + torch::Tensor reduce_weight; + torch::Tensor expert_id; + { + xllm::kernel::MoeFusedTopkParams moe_active_topk_params; + moe_active_topk_params.input = router_logits_2d; + moe_active_topk_params.topk = topk_; + moe_active_topk_params.num_expert_group = num_expert_group_; + moe_active_topk_params.topk_group = topk_group_; + moe_active_topk_params.normalize = renormalize_; + moe_active_topk_params.normed_by = "topk_logit"; + moe_active_topk_params.scoring_func = scoring_func_; + moe_active_topk_params.route_scale = route_scale_; + moe_active_topk_params.e_score_correction_bias = e_score_correction_bias; + std::tie(reduce_weight, expert_id) = + xllm::kernel::moe_active_topk(moe_active_topk_params); + } + + // Step 2: generate expert ids + torch::Tensor gather_idx; + torch::Tensor combine_idx; + torch::Tensor token_count; + std::optional cusum_token_count; + { + xllm::kernel::MoeGenIdxParams moe_gen_idx_params; + moe_gen_idx_params.expert_id = expert_id; + moe_gen_idx_params.expert_num = num_total_experts_; + std::vector output_vec = + xllm::kernel::moe_gen_idx(moe_gen_idx_params); + gather_idx = output_vec[0]; + combine_idx = output_vec[1]; + token_count = output_vec[2]; + // during all2all communication, we do not need cusum_token_count in the + // following computation + if (enable_all2all_communication) { + cusum_token_count = std::nullopt; + } else { + cusum_token_count = output_vec[3]; + } + } + + // Step 3: expand and quantize input if needed + torch::Tensor expand_hidden_states; + torch::Tensor hidden_states_scale; + torch::Tensor token_count_slice; + // all2all related variables + torch::Tensor dispatch_send_token_tensor; + // in all2all, the input is scattered, so there is no need to slice the token + // count, and we can use the dispatch buffer directly + if (enable_all2all_communication) { + token_count_slice = token_count; + int64_t num_token_expand = hidden_states_2d.size(0) * topk_; + int64_t dispatch_bytes = + num_token_expand * deep_ep_params_.dispatch_token_size; + dispatch_send_token_tensor = + deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes) + .view({num_token_expand, deep_ep_params_.dispatch_token_size}); + } else { + token_count_slice = + token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size); + } + + if (is_smoothquant_) { + xllm::kernel::ScaledQuantizeParams scaled_quantize_params; + scaled_quantize_params.x = hidden_states_2d; + // use dispatch_send_token_tensor buffer for input + // to reduce memory footprint + if (enable_all2all_communication) { + scaled_quantize_params.smooth = input_smooth_; + scaled_quantize_params.output = + dispatch_send_token_tensor.slice(1, 0, hidden_size_); + } else { + scaled_quantize_params.smooth = input_smooth_.slice( + 0, start_expert_id_, start_expert_id_ + expert_size); + scaled_quantize_params.gather_index_start_position = + cusum_token_count.value().index({start_expert_id_}).unsqueeze(0); + } + scaled_quantize_params.token_count = token_count_slice; + scaled_quantize_params.gather_index = gather_idx; + scaled_quantize_params.act_mode = "none"; + scaled_quantize_params.active_coef = 1.0; + scaled_quantize_params.is_gated = false; + scaled_quantize_params.quant_type = torch::kChar; + std::tie(expand_hidden_states, hidden_states_scale) = + xllm::kernel::scaled_quantize(scaled_quantize_params); + if (enable_all2all_communication) { + // since view_as_dtype has not supported stride yet, + // we need to copy the scale output to the dispatch buffer + torch::Tensor dispatch_scale_slice = + dispatch_send_token_tensor.slice(1, hidden_size_); + torch::Tensor hidden_states_scale_bytes = + view_as_dtype(hidden_states_scale, torch::kInt8) + .view_as(dispatch_scale_slice); + dispatch_scale_slice.copy_(hidden_states_scale_bytes); + } + } else { + xllm::kernel::MoeExpandInputParams moe_expand_input_params; + moe_expand_input_params.input = hidden_states_2d; + moe_expand_input_params.gather_index = gather_idx; + moe_expand_input_params.combine_idx = combine_idx; + moe_expand_input_params.topk = topk_; + expand_hidden_states = + xllm::kernel::moe_expand_input(moe_expand_input_params); + if (enable_all2all_communication) { + // use copy to place the output inside the dispatch buffer + torch::Tensor dispatch_tensor = + view_as_dtype(expand_hidden_states, torch::kChar); + dispatch_send_token_tensor.copy_(dispatch_tensor); + } + } + + // collect the selected tensor + selected_expert_info.reduce_weight = reduce_weight; + selected_expert_info.combine_idx = combine_idx; + selected_expert_info.token_count_slice = token_count_slice; + selected_expert_info.cusum_token_count = cusum_token_count; + if (is_smoothquant_) { + selected_expert_info.input_scale = hidden_states_scale; + } + + return expand_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication) { + if (!stream_initialized_) { + // update device record + device_ = xllm::Device(hidden_states.device()); + + // acquire streams from the pool again + routed_stream_ = device_.get_stream_from_pool(); + shared_stream_ = device_.get_stream_from_pool(); + stream_initialized_ = true; + } + + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + + // prepare the parameters for MoE computation + torch::Tensor shared_expert_output; + torch::IntArrayRef hidden_states_shape = hidden_states.sizes(); + torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType(); + torch::Tensor hidden_states_2d = + hidden_states.reshape({-1, hidden_states.size(-1)}); + torch::Tensor router_logits_2d = + router_logits.reshape({-1, router_logits.size(-1)}); + int64_t group_gemm_max_dim = enable_all2all_communication + ? deep_ep_params_.max_num_tokens_recv / topk_ + : hidden_states_2d.size(0); + int64_t expert_size = w13_.size(0); + + // Step 1-3: select experts + SelectedExpertInfo selected_expert_info; + torch::Tensor expand_hidden_states = + select_experts(hidden_states_2d, + router_logits_2d, + selected_expert_info, + enable_all2all_communication); + + // Communciation Step 1: Dipatch + // intermediate outputs that are used both in dispatch and combine + torch::Tensor gather_by_rank_index; + torch::Tensor token_sum; + if (enable_all2all_communication) { + int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_; + + // 1. Dispatch Step: Generate layout and send data + deep_ep_->dispatch_step(dispatch_token_num, + selected_expert_info.token_count_slice); + + // 2. Process Result: Generate indices and unpack to computation buffer + // use the buffer during initialization for the output + expand_hidden_states = dispatch_recv_token_tensor_head_; + std::optional output_tail = std::nullopt; + if (is_smoothquant_) { + output_tail = dispatch_recv_token_tensor_tail_; + // update selected_expert_info with the tail (input scale) + selected_expert_info.input_scale = output_tail; + } + + DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result( + num_experts_per_rank_, expand_hidden_states, output_tail); + + // Extract metadata for subsequent steps + gather_by_rank_index = deep_ep_meta.gather_rank_index; + selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice; + token_sum = deep_ep_meta.token_sum; + } + + // common gemm workspace for reduce memory footprint + torch::Tensor gemm_workspace; + + // Step 4: group gemm 1 + torch::Tensor gemm1_out = + create_group_gemm_output(expand_hidden_states, + w13_, + selected_expert_info.token_count_slice, + hidden_states_dtype, + gemm_workspace); + // ensure the lifespan of these parameters via brace + { + xllm::kernel::GroupGemmParams group_gemm_params; + torch::ScalarType a_dtype = + is_smoothquant_ ? torch::kInt8 : hidden_states_dtype; + group_gemm_params.a = + view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_}); + group_gemm_params.b = w13_; + group_gemm_params.token_count = + selected_expert_info.token_count_slice.to("cpu"); + if (is_smoothquant_) { + torch::Tensor a_scale = + selected_expert_info.input_scale.value().flatten(); + selected_expert_info.input_scale = + view_as_dtype(a_scale, torch::kFloat32); + group_gemm_params.a_scale = selected_expert_info.input_scale; + group_gemm_params.b_scale = w13_scale_; + } + group_gemm_params.max_dim = group_gemm_max_dim; + group_gemm_params.trans_a = false; + group_gemm_params.trans_b = true; + group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1; + group_gemm_params.output = gemm1_out; + group_gemm_params.combine_idx = std::nullopt; + gemm1_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Step 5: activation or scaled quantization(fused with activation) + torch::Tensor act_out; + torch::Tensor act_out_scale; + if (is_smoothquant_) { + int64_t slice_dim = gemm1_out.size(1); + if (is_gated_) slice_dim /= 2; + // slice operation is a view, does not take up extra memory, but points to + // the same memory + act_out = expand_hidden_states.slice(1, 0, slice_dim); + act_out_scale = + selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0)); + // call scaled quantization kernel (also fused with activation) + xllm::kernel::ScaledQuantizeParams scaled_quantize_params; + scaled_quantize_params.x = gemm1_out; + scaled_quantize_params.smooth = act_smooth_; + scaled_quantize_params.token_count = selected_expert_info.token_count_slice; + scaled_quantize_params.output = act_out; + scaled_quantize_params.output_scale = act_out_scale; + scaled_quantize_params.act_mode = hidden_act_; + scaled_quantize_params.active_coef = 1.0; + scaled_quantize_params.is_gated = is_gated_; + scaled_quantize_params.quant_type = torch::kChar; + std::tie(act_out, act_out_scale) = + xllm::kernel::scaled_quantize(scaled_quantize_params); + } else { + act_out = is_gated_ + ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous() + : gemm1_out; + // call activation kernel + xllm::kernel::ActivationParams activation_params; + activation_params.input = gemm1_out; + activation_params.output = act_out; + activation_params.cusum_token_count = + selected_expert_info.cusum_token_count; + activation_params.act_mode = hidden_act_; + activation_params.is_gated = is_gated_; + activation_params.start_expert_id = start_expert_id_; + activation_params.expert_size = expert_size; + xllm::kernel::active(activation_params); + } + + // Step 6: group gemm 2 + torch::Tensor gemm2_out = + create_group_gemm_output(act_out, + w2_, + selected_expert_info.token_count_slice, + hidden_states_dtype, + gemm_workspace); + // ensure the lifespan of these parameters via brace + { + xllm::kernel::GroupGemmParams group_gemm_params; + group_gemm_params.a = act_out; + group_gemm_params.b = w2_; + group_gemm_params.token_count = + selected_expert_info.token_count_slice.to("cpu"); + if (is_smoothquant_) { + group_gemm_params.a_scale = act_out_scale; + group_gemm_params.b_scale = w2_scale_; + } + group_gemm_params.max_dim = group_gemm_max_dim; + group_gemm_params.trans_a = false; + group_gemm_params.trans_b = true; + group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1; + group_gemm_params.output = gemm2_out; + group_gemm_params.combine_idx = selected_expert_info.combine_idx; + gemm2_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Communciation Step 2: Combine + if (enable_all2all_communication) { + int64_t num_token_expand = hidden_states_2d.size(0) * topk_; + // Delegate pack, layout generation and combine to DeepEP + torch::Tensor combine_send_layout = + deep_ep_->combine_step_pack(gemm2_out, + gather_by_rank_index, + token_sum, + hidden_size_, + hidden_states_dtype); + + // create a wait event for the current stream to finish computation + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + // pure communciation kernel: dispatch + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + gemm2_out = deep_ep_->combine_step_comm(combine_send_layout, + num_token_expand, + hidden_size_, + hidden_states_dtype); + } + + // pure computation kernel: shared experts + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + shared_expert_output = shared_experts_(hidden_states); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + } + } + + // After group gemm is finished, some tensors are no + // longer needed. We must explicitly release the memory. + expand_hidden_states = torch::Tensor(); + selected_expert_info.input_scale = std::nullopt; + act_out = torch::Tensor(); + + // Step 7: combine the intermediate results and get the final hidden states + torch::Tensor final_hidden_states; + // ensure the lifespan of these parameters via brace + { + xllm::kernel::MoeCombineResultParams moe_combine_result_params; + moe_combine_result_params.input = gemm2_out; + moe_combine_result_params.reduce_weight = + selected_expert_info.reduce_weight; + moe_combine_result_params.gather_ids = selected_expert_info.combine_idx; + moe_combine_result_params.cusum_token_count = + selected_expert_info.cusum_token_count; + moe_combine_result_params.start_expert_id = start_expert_id_; + moe_combine_result_params.expert_size = expert_size; + moe_combine_result_params.bias = std::nullopt; + // if all2all communication is enabled and shared output is provided, + // we will fused the add up to combine result + if (enable_all2all_communication && n_shared_experts_ > 0) { + moe_combine_result_params.residual = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + final_hidden_states = + xllm::kernel::moe_combine_result(moe_combine_result_params); + } + + // reshape the final hidden states to the original shape + final_hidden_states = final_hidden_states.reshape(hidden_states_shape); + + if (enable_all2all_communication) { + return final_hidden_states; + } + + // Communciation Step 3: AllReduce for non-all2all communication + // shared experts can be parallelized with the final communication step + // during moe computation. + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + if (tp_pg_->world_size() > 1) { + final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_); + } + if (parallel_args_.ep_size() > 1) { + final_hidden_states = parallel_state::reduce( + final_hidden_states, parallel_args_.moe_ep_group_); + } + } + + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + // for non all2all, we compute the shared experts parallelized with the + // final communication step + shared_expert_output = shared_experts_(hidden_states); + shared_expert_output = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + final_hidden_states += shared_expert_output; + } + + return final_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params) { + // we only support all2all communication for decode stage for now + bool enable_all2all_communication = + enable_deep_ep_ && std::all_of(input_params.parallel.dp_is_decode.begin(), + input_params.parallel.dp_is_decode.end(), + [](int32_t val) { return val == 1; }); + + bool is_dp_ep_parallel = + parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1; + // during all2all communication, the output has been + // gathered and sliced by dispatch and combine steps, + // so we do not need to gather input and slice output again + bool need_gather_and_slice = + is_dp_ep_parallel && !enable_all2all_communication; + + auto input = hidden_states; + if (need_gather_and_slice) { + input = parallel_state::gather(input, + parallel_args_.dp_local_process_group_, + input_params.parallel.dp_global_token_nums); + } + // MoE Gate + auto router_logits = gate_(input); + + // MoE Experts + auto output = + forward_experts(input, router_logits, enable_all2all_communication); + + if (need_gather_and_slice) { + output = get_dp_local_slice(output, input_params, parallel_args_); + } + + return output; +} + +void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) { + if (e_score_correction_bias_.defined() && + !e_score_correction_bias_is_loaded_) { + LOAD_WEIGHT(e_score_correction_bias); + } +} + +void FusedMoEImpl::load_experts(const StateDict& state_dict) { + const int64_t rank = tp_pg_->rank(); + const int64_t world_size = tp_pg_->world_size(); + const int64_t start_expert_id = start_expert_id_; + const int64_t num_experts_per_rank = num_experts_per_rank_; + const int64_t num_total_experts = num_total_experts_; + std::vector prefixes = {"gate_proj.", "up_proj."}; + if (is_smoothquant_) { + LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13); + LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale); + // When supporting DeepEP All2All mode, + // we need to load the complete set of expert weights corresponding to + // "up_proj.smooth". Note that even if deep EP mode is not enabled, it + // remains possible to retrieve the smooth quantization information for a + // subset of experts. Therefore, we intentionally do not check whether + // deep_ep_ is enabled in this case. + LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1); + LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1); + LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1); + LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0); + } else { + LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13); + LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1); + } +} + +void FusedMoEImpl::load_state_dict(const StateDict& state_dict) { + if (state_dict.size() == 0) { + return; + } + + if (n_shared_experts_ > 0) { + shared_experts_->load_state_dict( + state_dict.get_dict_with_prefix("shared_experts.")); + } + gate_->load_state_dict(state_dict.get_dict_with_prefix("gate.")); + load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate.")); + load_experts(state_dict.get_dict_with_prefix("experts.")); +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h new file mode 100644 index 00000000..8d4e9da9 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h @@ -0,0 +1,131 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "layers/common/deep_ep.h" +#include "layers/common/dense_mlp.h" +#include "layers/common/fused_moe_base.h" +#include "layers/common/linear.h" +#include "platform/device.h" +#include "util/tensor_helper.h" + +namespace xllm { +namespace layer { + +class FusedMoEImpl : public torch::nn::Module { + public: + FusedMoEImpl() = default; + FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + torch::Tensor forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication); + torch::Tensor forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params); + void load_state_dict(const StateDict& state_dict); + + private: + // struct to store the selected expert info + struct SelectedExpertInfo { + torch::Tensor reduce_weight; + torch::Tensor combine_idx; + torch::Tensor token_count_slice; + std::optional cusum_token_count; + std::optional input_scale; + }; + + // initial steps for MoE computation, select the experts for each token + torch::Tensor select_experts(const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info, + bool enable_all2all_communication); + + private: + int64_t num_total_experts_; + int64_t topk_; + int64_t num_expert_group_; + int64_t topk_group_; + double route_scale_; + int64_t hidden_size_; + int64_t n_shared_experts_; + bool is_gated_; + int64_t renormalize_; + std::string hidden_act_; + std::string scoring_func_; + bool is_smoothquant_; + + int64_t num_experts_per_rank_; + int64_t start_expert_id_; + + // Deep EP related parameters + bool enable_deep_ep_; + DeepEPBuffer deep_ep_buffer_; + DeepEPParams deep_ep_params_; + torch::Tensor dispatch_recv_token_tensor_head_; + torch::Tensor dispatch_recv_token_tensor_tail_; + + // steams for parallel shared experts + std::unique_ptr shared_stream_; + std::unique_ptr routed_stream_; + xllm::Device device_; + bool stream_initialized_ = false; + + ReplicatedLinear gate_{nullptr}; + DenseMLP shared_experts_{nullptr}; + DeepEP deep_ep_{nullptr}; + + QuantArgs quant_args_; + ParallelArgs parallel_args_; + torch::TensorOptions options_; + ProcessGroup* tp_pg_; + + DEFINE_WEIGHT(w13); + DEFINE_FUSED_WEIGHT(w1); + DEFINE_FUSED_WEIGHT(w3); + DEFINE_FUSED_WEIGHT(w2); + DEFINE_WEIGHT(e_score_correction_bias); + DEFINE_WEIGHT(w13_scale); + DEFINE_FUSED_WEIGHT(w1_scale); + DEFINE_FUSED_WEIGHT(w3_scale); + DEFINE_FUSED_WEIGHT(w2_scale); + DEFINE_FUSED_WEIGHT(input_smooth); + DEFINE_FUSED_WEIGHT(act_smooth); + + void load_e_score_correction_bias(const StateDict& state_dict); + void load_experts(const StateDict& state_dict); + // create the group gemm output tensor with the workspace + torch::Tensor create_group_gemm_output(const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype, + torch::Tensor& workspace); +}; +TORCH_MODULE(FusedMoE); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp new file mode 100644 index 00000000..0b97be6d --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp @@ -0,0 +1,219 @@ +/* Copyright 2025-2026 The xLLM Authors. +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 + https://github.com/jd-opensource/xllm/blob/main/LICENSE +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. +==============================================================================*/ + +#include "qwen3_5_gated_delta_net.h" + +#include + +namespace xllm { +namespace layer { + +Qwen3_5GatedDeltaNetImpl::Qwen3_5GatedDeltaNetImpl( + const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : Qwen3NextGatedDeltaNetImpl(args, + quant_args, + parallel_args, + options, + /*init_projections=*/false) { + in_proj_qkv_ = register_module("in_proj_qkv", + ColumnParallelLinear(args.hidden_size(), + k_size_ * 2 + v_size_, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); + in_proj_z_ = register_module("in_proj_z", + ColumnParallelLinear(args.hidden_size(), + v_size_, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); + in_proj_b_ = register_module("in_proj_b", + ColumnParallelLinear(args.hidden_size(), + num_v_heads_, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); + in_proj_a_ = register_module("in_proj_a", + ColumnParallelLinear(args.hidden_size(), + num_v_heads_, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); +} + +torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_qkvz_from_split_activations( + const torch::Tensor& qkv, + const torch::Tensor& z) const { + CHECK_EQ(qkv.dim(), 3) << "Expected qkv activation to be 3D, got " + << qkv.sizes(); + CHECK_EQ(z.dim(), 3) << "Expected z activation to be 3D, got " << z.sizes(); + CHECK_EQ(qkv.size(0), z.size(0)) << "qkv/z batch size mismatch."; + CHECK_EQ(qkv.size(1), z.size(1)) << "qkv/z sequence size mismatch."; + CHECK_EQ(qkv.size(2), (2 * k_size_ + v_size_) / tp_size_) + << "Unexpected qkv hidden size for Qwen3.5."; + CHECK_EQ(z.size(2), v_size_ / tp_size_) + << "Unexpected z hidden size for Qwen3.5."; + CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive."; + CHECK_EQ(num_v_heads_ % num_k_heads_, 0) + << "linear_num_value_heads must be divisible by linear_num_key_heads."; + + const int64_t bs = qkv.size(0); + const int64_t seqlen = qkv.size(1); + const int64_t local_k_heads = num_k_heads_ / tp_size_; + const int64_t local_v_heads = num_v_heads_ / tp_size_; + const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_; + + auto qkv_split = torch::split( + qkv, {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}, 2); + auto q = qkv_split[0].view({bs, seqlen, local_k_heads, head_k_dim_}); + auto k = qkv_split[1].view({bs, seqlen, local_k_heads, head_k_dim_}); + auto v = qkv_split[2].view({bs, seqlen, local_v_heads, head_v_dim_}); + auto z_view = z.view({bs, seqlen, local_v_heads, head_v_dim_}); + + v = v.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_}); + z_view = + z_view.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_}); + + return torch::cat({q, k, v, z_view}, -1).view({bs, seqlen, -1}).contiguous(); +} + +torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations( + const torch::Tensor& b, + const torch::Tensor& a) const { + CHECK_EQ(b.dim(), 3) << "Expected b activation to be 3D, got " << b.sizes(); + CHECK_EQ(a.dim(), 3) << "Expected a activation to be 3D, got " << a.sizes(); + CHECK_EQ(b.size(0), a.size(0)) << "b/a batch size mismatch."; + CHECK_EQ(b.size(1), a.size(1)) << "b/a sequence size mismatch."; + CHECK_EQ(b.size(2), num_v_heads_ / tp_size_) + << "Unexpected b hidden size for Qwen3.5."; + CHECK_EQ(a.size(2), num_v_heads_ / tp_size_) + << "Unexpected a hidden size for Qwen3.5."; + CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive."; + CHECK_EQ(num_v_heads_ % num_k_heads_, 0) + << "linear_num_value_heads must be divisible by linear_num_key_heads."; + + const int64_t bs = b.size(0); + const int64_t seqlen = b.size(1); + const int64_t local_k_heads = num_k_heads_ / tp_size_; + const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_; + + auto b_view = b.view({bs, seqlen, local_k_heads, num_v_heads_per_k}); + auto a_view = a.view({bs, seqlen, local_k_heads, num_v_heads_per_k}); + return torch::cat({b_view, a_view}, -1).view({bs, seqlen, -1}).contiguous(); +} + +std::pair +Qwen3_5GatedDeltaNetImpl::project_decode_inputs( + const torch::Tensor& hidden_states) { + const auto reshape_projection = [](const torch::Tensor& projection) { + return projection.view({projection.size(0), -1, projection.size(-1)}); + }; + auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states)); + auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states)); + auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states)); + auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states)); + return {merge_qkvz_from_split_activations(qkv, z_proj), + merge_ba_from_split_activations(b_proj, a_proj)}; +} + +std::pair +Qwen3_5GatedDeltaNetImpl::project_flat_inputs( + const torch::Tensor& hidden_states) { + auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0); + auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0); + auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0); + auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0); + auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj); + auto ba = merge_ba_from_split_activations(b_proj, a_proj); + return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(), + ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()}; +} + +std::optional< + std::tuple> +Qwen3_5GatedDeltaNetImpl::project_split_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) { + auto qkv = reshape_projected_tokens_with_pad( + attn_metadata, in_proj_qkv_->forward(hidden_states)); + auto z_proj = reshape_projected_tokens_with_pad( + attn_metadata, in_proj_z_->forward(hidden_states)); + auto b_proj = reshape_projected_tokens_with_pad( + attn_metadata, in_proj_b_->forward(hidden_states)); + auto a_proj = reshape_projected_tokens_with_pad( + attn_metadata, in_proj_a_->forward(hidden_states)); + + const int64_t batch_size = qkv.size(0); + const int64_t seq_len = qkv.size(1); + auto z = + z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); + auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_}); + auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_}); + return std::make_tuple(qkv, z, b, a); +} + +void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict( + const StateDict& state_dict) { + auto in_proj_qkv_state_dict = state_dict.get_dict_with_prefix("in_proj_qkv."); + if (in_proj_qkv_state_dict.size() > 0 && !in_proj_qkv_->is_weight_loaded()) { + in_proj_qkv_->load_state_dict( + in_proj_qkv_state_dict, + /*shard_tensor_count=*/3, + /*shard_sizes=*/ + {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}); + } + + auto in_proj_z_state_dict = state_dict.get_dict_with_prefix("in_proj_z."); + if (in_proj_z_state_dict.size() > 0 && !in_proj_z_->is_weight_loaded()) { + in_proj_z_->load_state_dict(in_proj_z_state_dict); + } + + auto in_proj_b_state_dict = state_dict.get_dict_with_prefix("in_proj_b."); + if (in_proj_b_state_dict.size() > 0 && !in_proj_b_->is_weight_loaded()) { + in_proj_b_->load_state_dict(in_proj_b_state_dict); + } + + auto in_proj_a_state_dict = state_dict.get_dict_with_prefix("in_proj_a."); + if (in_proj_a_state_dict.size() > 0 && !in_proj_a_->is_weight_loaded()) { + in_proj_a_->load_state_dict(in_proj_a_state_dict); + } +} + +void Qwen3_5GatedDeltaNetImpl::verify_projection_weights( + const std::string& prefix) const { + CHECK(in_proj_qkv_ && in_proj_qkv_->is_weight_loaded()) + << "Missing required weight after all shards loaded: " << prefix + << "in_proj_qkv.weight"; + CHECK(in_proj_z_ && in_proj_z_->is_weight_loaded()) + << "Missing required weight after all shards loaded: " << prefix + << "in_proj_z.weight"; + CHECK(in_proj_b_ && in_proj_b_->is_weight_loaded()) + << "Missing required weight after all shards loaded: " << prefix + << "in_proj_b.weight"; + CHECK(in_proj_a_ && in_proj_a_->is_weight_loaded()) + << "Missing required weight after all shards loaded: " << prefix + << "in_proj_a.weight"; +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h new file mode 100644 index 00000000..7c782e3f --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h @@ -0,0 +1,66 @@ +/* Copyright 2025-2026 The xLLM Authors. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include + +#include "qwen3_next_gated_delta_net.h" + +namespace xllm { +namespace layer { + +class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl { + public: + Qwen3_5GatedDeltaNetImpl() = default; + Qwen3_5GatedDeltaNetImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + protected: + std::pair project_decode_inputs( + const torch::Tensor& hidden_states) override; + std::pair project_flat_inputs( + const torch::Tensor& hidden_states) override; + std::optional< + std::tuple> + project_split_inputs(const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) override; + bool use_fla_ssm_state_layout() const override { return true; } + + void load_projection_state_dict(const StateDict& state_dict) override; + void verify_projection_weights(const std::string& prefix) const override; + + private: + torch::Tensor merge_qkvz_from_split_activations(const torch::Tensor& qkv, + const torch::Tensor& z) const; + torch::Tensor merge_ba_from_split_activations(const torch::Tensor& b, + const torch::Tensor& a) const; + + ColumnParallelLinear in_proj_qkv_{nullptr}; + ColumnParallelLinear in_proj_z_{nullptr}; + ColumnParallelLinear in_proj_b_{nullptr}; + ColumnParallelLinear in_proj_a_{nullptr}; +}; +TORCH_MODULE(Qwen3_5GatedDeltaNet); + +} // namespace layer +} // namespace xllm