diff --git "a/runtime-v2/overlay/python/sglang/srt/configs/model_config.py" "b/runtime-v2/overlay/python/sglang/srt/configs/model_config.py" new file mode 100644--- /dev/null +++ "b/runtime-v2/overlay/python/sglang/srt/configs/model_config.py" @@ -0,0 +1,2875 @@ +# Copyright 2023-2024 SGLang Team +# 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. +# ============================================================================== + +import copy +import json +import logging +import math +import os +from enum import Enum, IntEnum, auto +from functools import cached_property, lru_cache +from pathlib import Path +from typing import Any, Callable, List, Optional, Set, Union + +import torch +from transformers import PretrainedConfig +from transformers.utils import cached_file + +from sglang.srt.arg_groups.overrides import resolving_view +from sglang.srt.configs.bailing_hybrid import is_bailing_multi_gate_enabled +from sglang.srt.configs.embedding_model_spec import resolve_embedding_model_spec +from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_config +from sglang.srt.environ import envs +from sglang.srt.layers.quantization import QUANTIZATION_METHODS +from sglang.srt.runtime_context import get_platform +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import is_hip, retry +from sglang.srt.utils.hf_transformers.config import POINTER_ADAPTER_BASES +from sglang.srt.utils.hf_transformers_utils import ( + get_config, + get_context_length, + get_generation_config, + get_hf_text_config, + get_sparse_attention_config, +) +from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri +from sglang.utils import is_in_ci + +logger = logging.getLogger(__name__) + +_MODEL_CONFIG_FACTORIES: dict[type[ServerArgs], Callable[..., "ModelConfig"]] = {} + + +def register_model_config_factory( + server_args_type: type[ServerArgs], factory: Callable[..., "ModelConfig"] +) -> None: + """Register a factory with the same signature as ModelConfig.from_server_args. + + Register before resolution or model construction in each process. The nearest + registered type in the argument record's MRO wins, so registrations cover + subclasses too. Repeating the same registration is harmless; replacing a + different factory for the same type is an error. + """ + if not issubclass(server_args_type, ServerArgs): + raise TypeError("model-config factories require a ServerArgs subclass") + previous = _MODEL_CONFIG_FACTORIES.get(server_args_type) + if previous is not None and previous is not factory: + raise ValueError( + f"A model-config factory is already registered for {server_args_type.__qualname__}" + ) + _MODEL_CONFIG_FACTORIES[server_args_type] = factory + + +MIMO_V2_MODEL_ARCHS = ( + "MiMoV2ForCausalLM", + "MiMoV2FlashForCausalLM", +) +MIMO_V2_MULTIMODAL_ARCHS = ("MiMoV2ForCausalLM",) + +BAILING_MULTI_GATE_MM_ARCHS = frozenset( + { + "BailingMMNativeForConditionalGeneration", + "BailingMM2NativeForConditionalGeneration", + "BailingMoeV3VLForConditionalGeneration", + } +) + +SWA_SINK_ARCHS = frozenset( + { + "GptOssForCausalLM", + "GraniteSWAForCausalLM", + "GraniteMoeSWAForCausalLM", + "IQuestQ1ForCausalLM", + "IQuestQ1MTP", + } +) + + +def _quant_config_to_dict(quant_config): + if quant_config is not None and not isinstance(quant_config, dict): + return quant_config.to_dict() + return quant_config + + +def requires_mm_token_modalities( + model_architectures: Optional[List[str]], hf_text_config: PretrainedConfig +) -> bool: + """Whether a Bailing multimodal wrapper uses modality-specific routers.""" + return bool( + model_architectures + and any(arch in BAILING_MULTI_GATE_MM_ARCHS for arch in model_architectures) + and is_bailing_multi_gate_enabled(hf_text_config) + ) + + +def unwrap_modelopt_quantization_config(quant_config: dict) -> dict: + quantization = quant_config.get("quantization", quant_config) + if not isinstance(quantization, dict): + return {} + nested = quantization.get("quantization") + return nested if isinstance(nested, dict) else quantization + + +def get_mimo_v2_fused_qkv_expected_tp_size(hf_config): + layout = getattr(hf_config, "attention_projection_layout", None) + if layout in (None, "split"): + return None + if layout != "fused_qkv": + raise ValueError( + "MiMoV2 hf_config has unsupported " + f"attention_projection_layout={layout!r}; expected 'fused_qkv', " + "'split' or unset." + ) + + num_key_value_heads = getattr(hf_config, "num_key_value_heads", None) + text_config = getattr(hf_config, "text_config", None) + if num_key_value_heads is None and text_config is not None: + num_key_value_heads = getattr(text_config, "num_key_value_heads", None) + if num_key_value_heads is None: + raise ValueError( + "MiMoV2 hf_config has attention_projection_layout='fused_qkv' " + "but num_key_value_heads is missing; this value is required to " + "derive the fused qkv_proj TP size." + ) + return num_key_value_heads + + +class AttentionArch(IntEnum): + MLA = auto() + MHA = auto() + SSM = auto() # State Space Models (Mamba, Mamba2) + + +# Pure Mamba-1 (selective-scan) archs; same mixer/state layout, differing only +# in cosmetic details handled in their model files. +PURE_MAMBA1_ARCHITECTURES = ( + "FalconMambaForCausalLM", + "MambaForCausalLM", +) + +# Pure state-space (SSM) causal-LMs: no attention, so no num_attention_heads / +# head_dim in their HF config. Used for head-dim derivation and attention-arch +# detection below. +PURE_SSM_ARCHITECTURES = ("Mamba2ForCausalLM",) + PURE_MAMBA1_ARCHITECTURES + + +class ModelImpl(str, Enum): + AUTO = "auto" + SGLANG = "sglang" + TRANSFORMERS = "transformers" + MINDSPORE = "mindspore" + + +def _hf_arch(config) -> Optional[str]: + """First architecture from a HF config dict or PretrainedConfig (or None).""" + archs = ( + config.get("architectures") + if isinstance(config, dict) + else getattr(config, "architectures", None) + ) + return archs[0] if archs else None + + +def _hf_attr(config, name): + """Read an arbitrary field from a HF config dict or PretrainedConfig.""" + if isinstance(config, dict): + return config.get(name) + return getattr(config, name, None) + + +def is_deepseek_dsa(config) -> bool: + return ( + _hf_arch(config) + in ( + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + "DeepseekV3ForCausalLMNextN", + "MistralLarge3ForCausalLM", + "PixtralForConditionalGeneration", + "GlmMoeDsaForCausalLM", + "GlmMoeDsaForCausalLMNextN", + "Glm5NextForConditionalGenerationNextN", + "Glm5NextForConditionalGeneration", + "LongcatFlashForCausalLM", + "LongcatFlashForCausalLMNextN", + "Dots3NoteForCausalLM", + "Dots3NoteForCausalLMNextN", + "HYV4ForCausalLM", + "HYV4ForCausalLMNextN", + ) + and _hf_attr(config, "index_topk") is not None + ) + + +def dsa_draft_shards_under_dcp(config, dcp_size: int) -> bool: + """Whether a DSA EAGLE/NEXTN draft shards its KV like the DCP target. + + Only on ROCm; elsewhere the draft keeps a replicated pool and attends it + with dcp_size 1. + """ + return dcp_size > 1 and is_hip() and is_deepseek_dsa(config) + + +def is_glm_moe_dsa(config) -> bool: + """True for GLM-5.2, both the main arch and the NextN draft head.""" + return _hf_arch(config) in ( + "GlmMoeDsaForCausalLM", + "GlmMoeDsaForCausalLMNextN", + ) + + +def is_kimi_k3(config) -> bool: + return _hf_arch(config) in ( + "KimiK3ForConditionalGeneration", + "KimiK3LinearForCausalLM", + ) + + +def uses_kda_attention(config) -> bool: + configs = [config] + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + configs.append(get_text_config()) + else: + text_config = _hf_attr(config, "text_config") + if text_config is not None: + configs.append(text_config) + for config in configs: + linear_attn_config = _hf_attr(config, "linear_attn_config") + if isinstance(linear_attn_config, dict) and linear_attn_config.get( + "kda_layers" + ): + return True + layer_types = _hf_attr(config, "layer_types") or [] + if ( + "linear_attention" in layer_types + and _hf_attr(config, "linear_num_heads") is not None + and _hf_attr(config, "linear_head_dim") is not None + ): + return True + return False + + +def is_dspark_draft(config) -> bool: + return _hf_arch(config) in ("DSparkDraftModel", "DSparkMiniMaxDraftModel") + + +# The decision_config.json attention mode whose full-attention layers read the whole prompt. +NONCAUSAL_FULL_ATTENTION = "noncausal_full_attention" + + +@lru_cache +def load_decision_config(model_path: str, revision: Optional[str]) -> Optional[dict]: + """decision_config.json of a checkpoint whose LM head is a decision readout.""" + path = cached_file( + model_path, + "decision_config.json", + revision=revision, + _raise_exceptions_for_missing_entries=False, + ) + if path is None: + return None + with open(path) as f: + config = json.load(f) + if config["format_version"] != 1: + raise ValueError( + f"decision_config.json format_version {config['format_version']} " + "is not supported, only version 1 is" + ) + # Serving a mode this loader does not honor would silently change the answers. + attention_mode = config.get("attention_mode", "causal") + if attention_mode not in ("causal", NONCAUSAL_FULL_ATTENTION): + raise ValueError( + f"decision_config.json attention_mode {attention_mode!r} is not supported" + ) + pooling = config.get("pooling", "last") + if pooling != "last": + raise ValueError( + f"decision_config.json pooling {pooling!r} is not supported, only 'last' is" + ) + return config + + +@lru_cache +def load_joint_head_config(model_path: str, revision: Optional[str]) -> Optional[dict]: + """joint_head_config.json of a Clef checkpoint, whose joint schema head + answers /v1/systemone decisions from the backbone's hidden states.""" + path = cached_file( + model_path, + "joint_head_config.json", + revision=revision, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + if path is None: + return None + with open(path) as f: + return json.load(f) + + +# The values the Kev reference reads for head.pt fields that older checkpoints +# did not write. +_POINTER_HEAD_DEFAULTS = { + "base_revision": None, + "head_dim": 256, + "option_isolation": False, + "special_embeddings": False, + "weights_dtype": "fp32", + "temperature": 1.0, +} + + +@lru_cache +def load_pointer_head(model_path: str, revision: Optional[str]) -> Optional[dict]: + """head.pt of a Kev checkpoint, whose pointer head answers /v1/systemone + decisions from the backbone's hidden states.""" + path = cached_file( + model_path, + "head.pt", + revision=revision, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + if path is None: + return None + meta = { + **_POINTER_HEAD_DEFAULTS, + **torch.load(path, map_location="cpu", weights_only=True), + } + # Both change the prompt or the attention mask the head was trained with. + for name in ("option_isolation", "special_embeddings"): + if meta[name]: + raise ValueError(f"head.pt {name} is not supported") + return meta + + +@lru_cache +def load_pointer_adapter(model_path: str, revision: Optional[str]) -> Optional[dict]: + """head.pt of a Kev LoRA checkpoint with its merge scale. Such a checkpoint has + no model config and is served as the base backbone its head.pt pins, with the + adapter merged into the base weights at load.""" + meta = load_pointer_head(model_path, revision) + if meta is None: + return None + path = cached_file( + model_path, + "adapter_config.json", + revision=revision, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + if path is None: + return None + # Kev applies an adapter trained on a BF16 backbone unmerged, which a merge + # into the base weights does not reproduce. + if meta["weights_dtype"] != "fp32": + raise ValueError( + f"{model_path} is a Kev adapter trained on a {meta['weights_dtype']} " + "backbone, which SGLang does not serve. Serve its full weight release." + ) + with open(path) as f: + adapter = json.load(f) + return {**meta, "lora_scale": adapter["lora_alpha"] / adapter["r"]} + + +def is_qwen3_5(config) -> bool: + return _hf_arch(config) in ( + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForCausalLM", + "SenseNovaLoopForConditionalGeneration", + ) + + +def is_qwen3_5_mtp_draft(config) -> bool: + """The Qwen3.5 MoE MTP draft: _config_draft_model rewrites architectures[0] to + Qwen3_5ForCausalLMMTP before quantization is resolved.""" + return ( + _hf_arch(config) == "Qwen3_5ForCausalLMMTP" + and _hf_attr(config, "model_type") == "qwen3_5_moe" + ) + + +def is_deepseek_v4(config) -> bool: + return _hf_arch(config) in ( + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", + ) + + +def is_qwen4_exp(config) -> bool: + return _hf_arch(config) in ( + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpForCausalLMMTP", + ) + + +def resolve_spec_hidden_size( + hf_config, hidden_size: int, hc_mult: int +) -> tuple[int, Optional[int]]: + # DSV4 and Qwen4-Exp carry the hc-flattened stream across the target->draft + # boundary; other hc models (hy_v4) collapse to hidden_size first. + if hc_mult <= 1 or not (is_deepseek_v4(hf_config) or is_qwen4_exp(hf_config)): + return hidden_size, None + hc_hidden_size = hidden_size * hc_mult + return hc_hidden_size, hc_hidden_size + + +def get_dsa_index_head_dim(config: PretrainedConfig) -> int: + assert is_deepseek_dsa(config) or is_deepseek_v4(config) + return config.index_head_dim + + +def is_minimax_sparse(config: PretrainedConfig) -> bool: + arch = (config.architectures or [None])[0] + return arch in ( + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + ) + + +def uses_minimax_m3_nvfp4_native_stack(quantization: Optional[str]) -> bool: + """Whether an NVFP4 checkpoint runs its native Blackwell stack. + + The attention backend, the KV pool class and the KV cell-size budget must + all agree on this answer; sizing the pool for one layout while allocating + the other runs out of memory at startup. + """ + from sglang.srt.runtime_context import get_platform + + return quantization == "minimax_m3_nvfp4" and get_platform().is_blackwell + + +def get_minimax_sparse_attention_config(config: PretrainedConfig) -> dict: + text_cfg = getattr(config, "text_config", None) + cfg = ( + getattr(text_cfg, "sparse_attention_config", None) + if text_cfg is not None + else None + ) + if cfg is None: + cfg = getattr(config, "sparse_attention_config", None) + if cfg is None: + raise ValueError("Could not find sparse config. Is it MiniMax M3 Sparse model?") + return cfg + + +def get_minimax_sparse_layer_ids(sparse_cfg: dict) -> tuple[list[int], list[int]]: + sparse_freq = sparse_cfg["sparse_attention_freq"] + dense_layer_ids = [i for i, f in enumerate(sparse_freq) if f == 0] + sparse_layer_ids = [i for i, f in enumerate(sparse_freq) if f != 0] + return dense_layer_ids, sparse_layer_ids + + +def get_minimax_sparse_disable_value_layer_ids(sparse_cfg: dict) -> list[int]: + flags = sparse_cfg.get("sparse_disable_index_value") + if flags is None: + return [] + return [i for i, f in enumerate(flags) if f != 0] + + +def get_minimax_sparse_score_type(sparse_cfg: dict) -> str: + score_type = sparse_cfg.get("sparse_score_type", "max") + assert score_type in ( + "max", + "lse", + ), f"sparse_score_type must be 'max' or 'lse', got {score_type!r}" + return score_type + + +def get_dsa_index_topk(config: PretrainedConfig) -> int: + assert is_deepseek_dsa(config) + return config.index_topk + + +def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool: + """Return whether a DSA layer reuses the previous layer's top-k indices.""" + assert is_deepseek_dsa(config) + + indexer_types = getattr(config, "indexer_types", None) + if indexer_types is not None: + return ( + 0 <= layer_id < len(indexer_types) and indexer_types[layer_id] == "shared" + ) + + # LongCat computes fresh top-k indices every cli_factor layers. + cli_factor = getattr(config, "cli_factor", 1) + if cli_factor is None: + cli_factor = 1 + assert cli_factor > 0, f"cli_factor must be positive, got {cli_factor}" + if cli_factor > 1: + return layer_id % cli_factor != 0 + + pattern = getattr(config, "index_topk_pattern", None) + if pattern is not None: + return layer_id < len(pattern) and pattern[layer_id] == "S" + + freq = getattr(config, "index_topk_freq", 1) + if freq is None: + freq = 1 + assert freq > 0, f"index_topk_freq must be positive, got {freq}" + offset = getattr(config, "index_skip_topk_offset", None) + if offset is not None: + assert offset > 0, ( + "index_skip_topk_offset must be positive; offset <= 0 " + "marks layer 0 as skip_topk with no prior topk to reuse" + ) + return max(layer_id - offset + 1, 0) % freq != 0 + + return max(layer_id - 1, 0) % freq != 0 + + +def get_dsa_index_n_heads(config: PretrainedConfig) -> int: + assert is_deepseek_dsa(config) + return config.index_n_heads + + +def get_dsa_index_kpool(config: PretrainedConfig) -> int: + return getattr(config, "index_kpool", 1) + + +def get_dsa_mtp_topk_width(config: PretrainedConfig) -> int: + """MTP seeds include index_topk pooled tokens plus up to index_kpool - 1 tail tokens.""" + index_kpool = get_dsa_index_kpool(config) + assert index_kpool >= 1, f"index_kpool must be positive, got {index_kpool}" + return config.index_topk + index_kpool - 1 + + +def get_dsa_index_kpool_compress(config: PretrainedConfig) -> bool: + return getattr(config, "index_kpool_compress", False) + + +REQUANTIZATION_METHODS = ["quark_mxfp4"] + + +def get_num_indexer_layers(config) -> int: + """Layer count for the global indexer-topk capturer's host buffer. + + DSA models (V3.2) expose one capturer slot per transformer layer. With + index_topk_freq > 1 some layers reuse prev layer's topk; those still get a + slot mirrored at the MLA call site even if no Indexer module is built. + DSv4 has C4 indexers only on layers whose compress_ratio == 4. Other + architectures: set num_indexer_layers on hf_text_config; 0 disables the + capturer. + """ + if is_deepseek_dsa(config): + return config.num_hidden_layers + if is_deepseek_v4(config): + compress_ratios = getattr(config, "compress_ratios", None) or [] + return sum(1 for r in compress_ratios if r == 4) + return getattr(config, "num_indexer_layers", 0) + + +class ModelConfig: + def __init__( + self, + model_path: str, + trust_remote_code: bool = True, + revision: Optional[str] = None, + context_length: Optional[int] = None, + model_override_args: str = "{}", + is_embedding: Optional[bool] = None, + enable_multimodal: Optional[bool] = None, + dtype: str = "auto", + quantization: Optional[str] = None, + override_config_file: Optional[str] = None, + is_draft_model: bool = False, + model_impl: Union[str, ModelImpl] = ModelImpl.AUTO, + sampling_defaults: str = "openai", + quantize_and_serve: bool = False, + is_multi_layer_eagle: bool = False, + encoder_only: bool = False, + language_only: bool = False, + language_model_only: bool = False, + disable_hybrid_swa_memory: bool = False, + model_config_parser: str = "auto", + speculative_algorithm: Optional[str] = None, + is_draft_quantization_explicit: bool = False, + ) -> None: + # Parse args + self.model_path = model_path + self.revision = revision + self.quantization = quantization + self.is_draft_model = is_draft_model + self.is_draft_quantization_explicit = is_draft_quantization_explicit + self.speculative_algorithm = speculative_algorithm + self.model_impl = model_impl + self.sampling_defaults = sampling_defaults + self.quantize_and_serve = quantize_and_serve + self.is_multi_layer_eagle = is_multi_layer_eagle + self.disable_hybrid_swa_memory = disable_hybrid_swa_memory + self.model_config_parser = model_config_parser + + # Validate quantize_and_serve configuration + self._validate_quantize_and_serve_config() + + # Get hf config + self._maybe_pull_model_for_runai(self.model_path) + self._maybe_pull_model_tokenizer_from_remote() + self.model_override_args = json.loads(model_override_args) + kwargs = {} + if override_config_file and override_config_file.strip(): + kwargs["_configuration_file"] = override_config_file.strip() + # get_config() is cached. ModelConfig mutates hf_config for draft-model + # remapping and architecture-specific normalization, so each instance + # must own an isolated copy. + self.hf_config = copy.deepcopy( + get_config( + self.model_path, + trust_remote_code=trust_remote_code, + revision=revision, + model_override_args=self.model_override_args, + model_config_parser=model_config_parser, + **kwargs, + ) + ) + self.pointer_head = None + base = POINTER_ADAPTER_BASES.get((self.model_path, revision)) + if base is not None: + # A Kev LoRA checkpoint, whose base weights load with its adapter merged. + self.pointer_head = load_pointer_adapter(self.model_path, revision) + self.model_path, self.revision = base + revision = self.revision + self.hf_text_config = get_hf_text_config(self.hf_config) + self.decision_config = None + if self.hf_config.architectures == ["Qwen3_5Model"]: + self.decision_config = load_decision_config(self.model_path, revision) + if self.decision_config is not None: + # A bare backbone whose readout the loader places in the LM head. + self.hf_config.architectures = ["Qwen3_5ForConditionalGeneration"] + if ( + self.decision_config.get("attention_mode") + == NONCAUSAL_FULL_ATTENTION + ): + # Full attention sees the whole prompt, while linear attention stays causal. + self.hf_text_config.is_causal = False + self.joint_head_config = None + if self.hf_config.architectures == ["Qwen3_5ForConditionalGeneration"]: + self.joint_head_config = load_joint_head_config(self.model_path, revision) + if self.hf_config.architectures == ["Qwen3_5TextModel"]: + self.pointer_head = load_pointer_head(self.model_path, revision) + if self.pointer_head is not None: + # A Kev backbone without an LM head, served by the text-only class. + self.hf_config.architectures = ["Qwen3_5ForCausalLM"] + self.requires_mm_token_modalities = requires_mm_token_modalities( + self.hf_config.architectures, self.hf_text_config + ) + self.is_embedding_gemma = is_embedding_gemma(self.hf_text_config) + self.embedding_model_spec = resolve_embedding_model_spec( + self.hf_config.architectures, + is_embedding_requested=bool(is_embedding), + is_embedding_gemma=self.is_embedding_gemma, + ) + + rope_scaling = getattr(self.hf_text_config, "rope_parameters", None) or getattr( + self.hf_text_config, "rope_scaling", {} + ) + # Text-only serving comes from either the checkpoint's own declaration + # or the --language-model-only flag; every capability flag below + # (is_multimodal, is_*_understandable_model) must see both sources, + # since /model_info advertises them and drives media warmup requests. + self.is_lm_only = language_model_only or getattr( + self.hf_config, "language_model_only", False + ) + self.model_is_mrope = ( + not self.is_lm_only + and rope_scaling is not None + and "mrope_section" in rope_scaling + ) + + self.hf_generation_config = get_generation_config( + self.model_path, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + + # Set enable_multimodal + if enable_multimodal is None: + mm_disabled_models = [ + "Gemma3ForConditionalGeneration", + "Llama4ForConditionalGeneration", + "Step3VLForConditionalGeneration", + "InklingForConditionalGeneration", + ] + if ( + self.hf_config.architectures[0] in mm_disabled_models + and self.model_impl != ModelImpl.TRANSFORMERS + ): + enable_multimodal = False + logger.info( + f"Multimodal is disabled for {self.hf_config.model_type}. To enable it, set --enable-multimodal." + ) + elif self.hf_config.architectures[0] in MIMO_V2_MULTIMODAL_ARCHS and not ( + hasattr(self.hf_config, "vision_config") + and hasattr(self.hf_config, "audio_config") + ): + enable_multimodal = False + logger.info( + "Multimodal is disabled for this MiMoV2 checkpoint: " + "vision_config/audio_config not found in the model config " + "(likely a text-only MiMoV2 variant)." + ) + else: + enable_multimodal = True + + # Config draft model + self._config_draft_model() + + # Mixed FP8/MXFP4 ckpts mark mxfp4 routed experts via this key. + quantization_config = ( + _quant_config_to_dict(getattr(self.hf_config, "quantization_config", None)) + or {} + ) + self.hf_quant_config: dict = quantization_config + routed_experts_quant_method = quantization_config.get( + "routed_experts_quant_method" + ) + self.is_fp4_experts: bool = routed_experts_quant_method == "mxfp4" + if self.is_fp4_experts: + logger.info("Detected mixed checkpoint layout: routed experts are MXFP4.") + + # MiMo-V2 mxfp4 ckpts declare the routed-expert layout via store_dtype. + if ( + not self.is_fp4_experts + and _hf_arch(self.hf_config) in MIMO_V2_MODEL_ARCHS + and str(quantization_config.get("store_dtype") or "").lower() == "mxfp4" + ): + self.is_fp4_experts = True + logger.info("Detected MiMo-V2 mxfp4 routed-expert layout.") + + # DSV4 mxfp4 layout applies only when the ckpt does not opt in above. + if is_deepseek_v4(self.hf_config) and routed_experts_quant_method is None: + self.is_fp4_experts = envs.SGLANG_DSV4_FP4_EXPERTS.get() + if ( + not envs.SGLANG_DSV4_FP4_EXPERTS.is_set() + or envs.SGLANG_DSV4_FP4_DEQUANT.is_set() + ): + from sglang.srt.configs.deepseek_v4 import try_detect_fp4_experts + + detected = try_detect_fp4_experts(self.model_path) + if detected is not None: + self.is_fp4_experts = detected + logger.info( + "Auto-detected DSV4 routed-expert layout: is_fp4_experts=%s", + self.is_fp4_experts, + ) + if envs.SGLANG_DSV4_FP4_DEQUANT.get(): + envs.SGLANG_DSV4_FP4_DEQUANT.set(self.is_fp4_experts is not None) + + # HF config.json inherits topk_group=4 from the V3 template, but + # DSV4 trains with no group limiting (sqrtsoftplus + full-expert + # top-k). Force topk_group == n_group so deepseek_v2.py:531's + # `n_group > topk_group` evaluates False and routes to the + # ungrouped sqrtsoftplus path. The grouped impl only supports + # sigmoid scoring (topk.py:722) and would silently corrupt expert + # weights if hit. + n_group = getattr(self.hf_config, "n_group", None) + if n_group is not None: + self.hf_config.topk_group = n_group + + # Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4) + self.nvfp4_moe_meta: Optional[dict] = None + hybrid_quant_cfg = _quant_config_to_dict( + getattr(self.hf_config, "quantization_config", None) + ) + if ( + hybrid_quant_cfg is not None + and str(hybrid_quant_cfg.get("quant_algo", "")).upper() == "MIXED_PRECISION" + and str(hybrid_quant_cfg.get("moe_quant_algo", "")).upper() == "NVFP4" + and hybrid_quant_cfg.get("group_size") is not None + ): + self.nvfp4_moe_meta = { + "group_size": int(hybrid_quant_cfg["group_size"]), + "exclude_modules": list(hybrid_quant_cfg.get("ignore") or []), + } + logger.info( + "Auto-detected hybrid FP8+NVFP4 checkpoint " + "(NVFP4 MoE group_size=%d, %d exclude_modules)", + self.nvfp4_moe_meta["group_size"], + len(self.nvfp4_moe_meta["exclude_modules"]), + ) + + # Check model type + self.attention_chunk_size = getattr( + self.hf_text_config, "attention_chunk_size", None + ) + self.sliding_window_size = self._get_sliding_window_size() + self.is_generation = not self.is_embedding_gemma and is_generation_model( + self.hf_config.architectures, is_embedding + ) + # The vision_config/audio_config attribute heuristic is only applied when + # the transformers backend is explicitly requested. Some text-only models + # (e.g. xai-org/grok-2 with model_type="git") would otherwise be + # false-positively detected because their HF config auto-populates a + # `vision_config` in __post_init__. + has_multimodal_subconfig = self.hf_config is not self.hf_text_config or ( + self.model_impl == ModelImpl.TRANSFORMERS + and ( + hasattr(self.hf_config, "vision_config") + or hasattr(self.hf_config, "audio_config") + ) + ) + has_dsv41_vision = ( + self.hf_config.model_type == "deepseek_v41" + and self.hf_config.vision_n_layers > 0 + ) + self.is_multimodal = ( + enable_multimodal + and not self.is_lm_only + and ( + is_multimodal_model(self.hf_config.architectures) + or has_multimodal_subconfig + or has_dsv41_vision + ) + ) + self.is_audio_model = enable_multimodal and is_audio_model( + self.hf_config.architectures + ) + # Gated on `is_multimodal` because this flag is advertised via /model_info + # and drives the VLM warmup request, while the OpenAI serving layer rejects + # media input for models that are not `is_multimodal`. A text-only model + # with an auto-populated `vision_config` (see above) would otherwise warm up + # with an image request that its own serving layer answers with 400. + # Key on the tower, not the attribute: several config classes default + # vision_config to None, which presence alone would read as image-capable + # (MuseGlimmerConfig's text-only layouts are one such case). + # TODO: requires further polishing + self.is_image_understandable_model = ( + self.is_multimodal + and getattr(self.hf_config, "vision_config", None) is not None + ) + if self.is_multimodal and has_dsv41_vision: + self.is_image_understandable_model = True + + # Models expose audio_config at different nesting levels: + # - top-level audio_config: e.g. Qwen2Audio + # - thinker_config.audio_config: Qwen3-Omni, Qwen3-ASR (nested thinker arch) + # - sound_config: Nemotron AVLM with Parakeet audio encoder + # - is_audio_model(): Whisper, Qwen3-ASR (architecture-based fallback) + # TODO: Handle this more robustly by standardizing the config structure in the future + self.is_audio_understandable_model = ( + enable_multimodal + and not self.is_lm_only + and ( + hasattr(self.hf_config, "audio_config") + or hasattr( + getattr(self.hf_config, "thinker_config", None), "audio_config" + ) + or getattr(self.hf_config, "sound_config", None) is not None + or is_audio_model(self.hf_config.architectures) + ) + ) + + self.is_multimodal_chunked_prefill_supported = ( + enable_multimodal + and is_multimodal_chunked_prefill_supported(self.hf_config.architectures) + ) + self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures) + self.is_local_attention_model = is_local_attention_model( + self.hf_config.architectures + ) + self.use_ngram_embedding = getattr(self.hf_config, "use_ngram_embedding", False) + self.ngram_embedding_n = ( + self.hf_config.ngram_embedding_n if self.use_ngram_embedding else 0 + ) + self.use_engram = bool(getattr(self.hf_config, "engram_layer_ids", ())) + self._derive_multimodal_cuda_graph_support(enable_multimodal) + self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) + + # Derive context length and model shapes + self._derive_context_length(context_length) + self._derive_model_shapes() + + # Update hybrid model + self._derive_hybrid_model() + + # Verify quantization + self._verify_quantization() + + # Verify dual-chunk attention config + self._verify_dual_chunk_attention_config() + + # Cache attributes + self.hf_eos_token_id = self._get_hf_eos_token_id() + # Set by scheduler when reasoning_parser is enabled + self.think_end_ids: Optional[List[int]] = None + self.self_opened_think_start_ids: Optional[List[int]] = None + self.request_selectable_think_end_id_sequences: Optional[List[List[int]]] = None + + # multimodal + self.image_token_id = getattr( + self.hf_config, "image_token_id", None + ) or getattr(self.hf_config, "image_token_index", None) + + self.hf_config.encoder_only = encoder_only + self.hf_config.language_only = language_only + # Checkpoints declare this one themselves (hf_transformers/processor.py), + # so the flag may only turn it on: writing the default back would build a + # vision tower with no weights to fill. + self.hf_config.language_model_only = self.is_lm_only + + # matryoshka embeddings + self.matryoshka_dimensions = getattr( + self.hf_config, "matryoshka_dimensions", None + ) + self.is_matryoshka = self.matryoshka_dimensions or getattr( + self.hf_config, "is_matryoshka", False + ) + + @staticmethod + def from_server_args( + server_args: ServerArgs, + model_path: str = None, + model_revision: str = None, + is_draft_model: bool = False, + context_length: Optional[int] = None, + **kwargs, + ): + for record_type in type(server_args).__mro__: + factory = _MODEL_CONFIG_FACTORIES.get(record_type) + if factory is not None: + return factory( + server_args, + model_path=model_path, + model_revision=model_revision, + is_draft_model=is_draft_model, + context_length=context_length, + **kwargs, + ) + + cfg = resolving_view(server_args) + quantization = ( + cfg.speculative_draft_model_quantization + if is_draft_model + else cfg.quantization + ) + override_config_file = ( + cfg.decrypted_draft_config_file + if is_draft_model + else cfg.decrypted_config_file + ) + return ModelConfig( + model_path=model_path or cfg.model_path, + trust_remote_code=cfg.trust_remote_code, + revision=model_revision or cfg.revision, + context_length=( + context_length if context_length is not None else cfg.context_length + ), + model_override_args=cfg.json_model_override_args, + is_embedding=cfg.is_embedding, + enable_multimodal=cfg.enable_multimodal, + dtype=cfg.dtype, + quantization=quantization, + model_impl=cfg.model_impl, + sampling_defaults=cfg.sampling_defaults, + quantize_and_serve=cfg.quantize_and_serve, + override_config_file=override_config_file, + is_multi_layer_eagle=cfg.enable_multi_layer_eagle, + language_only=cfg.language_only, + language_model_only=cfg.language_model_only, + encoder_only=cfg.encoder_only, + is_draft_model=is_draft_model, + is_draft_quantization_explicit=( + is_draft_model and cfg._speculative_draft_quantization_explicitly_set + ), + disable_hybrid_swa_memory=cfg.disable_hybrid_swa_memory, + model_config_parser=cfg.model_config_parser, + speculative_algorithm=cfg.speculative_algorithm, + **kwargs, + ) + + def _derive_multimodal_cuda_graph_support(self, enable_multimodal: bool) -> None: + """Declare graph capabilities before deriving shapes and validating config. + + External ModelConfig subclasses can extend this without modifying global + architecture tables or repairing the config after construction. + """ + # A multimodal arch is piecewise-incompatible until its LM prefill is validated. + self.is_piecewise_cuda_graph_disabled_model = ( + is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures) + or ( + self.is_multimodal + and not is_multimodal_piecewise_cuda_graph_supported( + self.hf_config.architectures + ) + ) + ) + # The vision encoder still runs eagerly; this opt-in captures only the + # language-model prefill of architectures validated for piecewise graphs. + self.is_multimodal_piecewise_cuda_graph_supported = enable_multimodal and ( + is_multimodal_piecewise_cuda_graph_supported(self.hf_config.architectures) + ) + self.is_multimodal_breakable_cuda_graph_supported = enable_multimodal and ( + is_multimodal_breakable_cuda_graph_supported(self.hf_config.architectures) + ) + + def _config_draft_model(self): + is_draft_model = self.is_draft_model + + from sglang.srt.configs.dots3 import Dots3Config + + if is_draft_model and isinstance(self.hf_text_config, Dots3Config): + self.hf_config.architectures[0] = ( + self.hf_text_config.configure_draft_model() + ) + + if ( + is_draft_model + and self.speculative_algorithm == "DSPARK" + and self.hf_config.architectures[0] == "DSparkMiniMaxDraftModel" + ): + from sglang.srt.configs.minimax_vl import configure_minimax_m3_dspark + + self.hf_config = configure_minimax_m3_dspark(self.hf_config) + self.hf_text_config = self.hf_config + if self.is_draft_quantization_explicit and self.quantization is None: + self.hf_config.quantization_config = None + + if is_draft_model and self.hf_config.architectures[0] in [ + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + ]: + self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" + + if is_draft_model and self.hf_config.architectures[0] == "GlmMoeDsaForCausalLM": + self.hf_config.architectures[0] = "GlmMoeDsaForCausalLMNextN" + + if ( + is_draft_model + and self.hf_config.architectures[0] == "DeepseekV4ForCausalLM" + ): + from sglang.srt.speculative.dspark_components.dspark_config import ( + checkpoint_bundles_dspark_draft, + ) + + # A dspark-bundled checkpoint may also carry MTP layers; the + # selected algorithm decides which draft arch to load. + if checkpoint_bundles_dspark_draft(self.hf_config) and ( + self.speculative_algorithm in (None, "DSPARK") + ): + self.hf_config.architectures[0] = "DeepseekV4ForCausalLMDSpark" + logger.info( + "Draft checkpoint bundles a DSpark head; loading draft arch " + "DeepseekV4ForCausalLMDSpark." + ) + else: + self.hf_config.architectures[0] = "DeepseekV4ForCausalLMNextN" + self.hf_config.num_nextn_predict_layers = 1 + + if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM": + self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN" + if ( + is_draft_model + and self.hf_config.architectures[0] == "Glm4MoeLiteForCausalLM" + ): + self.hf_config.architectures[0] = "Glm4MoeLiteForCausalLMNextN" + + if ( + is_draft_model + and self.hf_config.architectures[0] == "Glm5NextForConditionalGeneration" + ): + self.hf_config.architectures[0] = "Glm5NextForConditionalGenerationNextN" + self.hf_text_config.architectures = list(self.hf_config.architectures) + self.hf_text_config.num_nextn_predict_layers = 1 + self.hf_text_config.linear_attn_config = None + + if is_draft_model and self.hf_config.architectures[0] in [ + "GlmOcrForConditionalGeneration", + ]: + self.hf_config.architectures[0] = "GlmOcrForConditionalGenerationNextN" + + if ( + is_draft_model + and self.hf_config.architectures[0] == "LongcatFlashForCausalLM" + ): + self.hf_config.architectures[0] = "LongcatFlashForCausalLMNextN" + self.hf_config.num_hidden_layers = self.hf_config.num_nextn_predict_layers + + if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM": + self.hf_config.architectures[0] = "MiMoMTP" + if is_draft_model and self.hf_config.architectures[0] in MIMO_V2_MODEL_ARCHS: + self.hf_config.architectures[0] = "MiMoV2MTP" + self.hf_config.num_nextn_predict_layers = 1 + if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM": + self.hf_config.architectures[0] = "Step3p5MTP" + if ( + is_draft_model + and self.hf_config.architectures[0] == "InklingForConditionalGeneration" + ): + self.hf_config.architectures[0] = "InklingForConditionalGenerationMTP" + if self.hf_text_config.mtp_local_layer_ids: + # A local MTP depth keeps the trunk's window unless it declares + # its own: the draft backend builds its window rail only if set. + self.sliding_window_size = ( + self.hf_text_config.mtp_local_extent or self.sliding_window_size + ) + if ( + is_draft_model + and self.hf_config.architectures[0] == "GigaChat35ForCausalLM" + ): + self.hf_config.architectures[0] = "GigaChat35ForCausalLMNextN" + if ( + is_draft_model + and self.hf_config.architectures[0] == "Step3p7ForConditionalGeneration" + ): + self.hf_config = self.hf_text_config + self.hf_config.architectures = ["Step3p5MTP"] + if is_draft_model and self.hf_config.architectures[0] in [ + "BailingMoeV2ForCausalLM", + "BailingMoeForCausalLM", + "BailingMoeV2_5ForCausalLM", + "BailingMoeV3ForCausalLM", + ]: + self.hf_config.architectures[0] = "BailingMoeForCausalLMNextN" + if ( + is_draft_model + and self.hf_config.architectures[0] == "Ernie4_5_MoeForCausalLM" + ): + self.hf_config.architectures[0] = "Ernie4_5_MoeForCausalLMMTP" + + if is_draft_model and self.hf_config.architectures[0] == "Qwen3NextForCausalLM": + self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + + if ( + is_draft_model + and self.hf_config.architectures[0] == "Qwen4ExpForConditionalGeneration" + ): + # The target's ModelConfig shares this hf_config object; deep-copy + # before the MTP rewrites below so the target keeps its full depth. + self.hf_config = copy.deepcopy(self.hf_config) + self.hf_text_config = get_hf_text_config(self.hf_config) + self.hf_config.architectures[0] = "Qwen4ExpForCausalLMMTP" + text_config = self.hf_text_config + text_config.num_nextn_predict_layers = 1 + # layers_block_type follows layer_types, not num_hidden_layers, + # so both must shrink for the draft's full_attention_layer_ids to be [0]. + text_config.num_hidden_layers = 1 + text_config.layer_types = ["full_attention"] + text_config.full_attention_interval = 1 + + if is_draft_model and self.hf_config.architectures[0] == "Qwen3MoeForCausalLM": + self.hf_config.architectures[0] = "Qwen3MoeForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + + if ( + is_draft_model + and self.hf_config.architectures[0] + == "SenseNovaLoopForConditionalGeneration" + ): + # MTP patch: the bundled mtp.* head is a plain Qwen3.5 MTP layer (one + # full-attention layer, BF16). Shrink the looped execution layout to + # that single layer; deep-copy so the target keeps its 56-slot plan. + self.hf_config = copy.deepcopy(self.hf_config) + self.hf_text_config = get_hf_text_config(self.hf_config) + text_config = self.hf_text_config + text_config.num_hidden_layers = 1 + text_config.layer_types = ["full_attention"] + text_config.full_attention_interval = 1 + text_config.execution_plan = (0,) + text_config.mhc = None + text_config.num_nextn_predict_layers = 1 + self.hf_config.num_nextn_predict_layers = 1 + self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP" + + if is_draft_model and self.hf_config.architectures[0] in [ + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForCausalLM", + "InternS2PreviewForConditionalGeneration", + "InternS2MobiusForConditionalGeneration", + ]: + if ( + self.hf_config.architectures[0] + == "InternS2MobiusForConditionalGeneration" + ): + # The target owns 2,560 experts through four shared physical + # banks, while its bundled MTP layer is an ordinary Qwen3.5 + # MoE layer with the checkpoint-declared smaller expert set. + self.hf_text_config.model_type = "qwen3_5_moe_text" + self.hf_text_config.num_experts = self.hf_text_config.mtp_num_experts + self.hf_text_config.num_experts_per_tok = ( + self.hf_text_config.mtp_num_experts_per_tok + ) + self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + self.hf_text_config.num_nextn_predict_layers = 1 + + if is_draft_model and self.hf_config.architectures[0] in ( + "ExaoneMoEForCausalLM", + "ExaoneMoeForCausalLM", + ): + self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + + if ( + is_draft_model + and self.hf_config.architectures[0] == "NemotronH_Omni_Reasoning_V3" + ): + self.hf_config = self.hf_text_config + self.hf_config.architectures = ["NemotronHForCausalLMMTP"] + self.hf_config.num_nextn_predict_layers = 1 + + if is_draft_model and self.hf_config.architectures[0] in [ + "NemotronHForCausalLM", + "NemotronHPuzzleForCausalLM", + ]: + self.hf_config.architectures[0] = "NemotronHForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + + if is_draft_model and self.hf_config.architectures[0] == "HYV3ForCausalLM": + self.hf_config.architectures[0] = "HYV3ForCausalLMNextN" + self.hf_config.num_nextn_predict_layers = 1 + + if is_draft_model and self.hf_config.architectures[0] == "HYV4ForCausalLM": + self.hf_config.architectures[0] = "HYV4ForCausalLMNextN" + self.hf_config.num_nextn_predict_layers = 1 + + def _derive_hybrid_model(self): + # Use self.context_len after it has been initialized to prevent using context_len which may be None. + self.is_hybrid_swa = ( + is_hybrid_swa_model(self.hf_config.architectures, self.hf_text_config) + and not self.disable_hybrid_swa_memory + ) + # Whole-model split, read-only; per-runner slices live on ModelLayerInfo. + self.swa_attention_layer_ids: Optional[List[int]] = None + self.full_attention_layer_ids: Optional[List[int]] = None + + if self.is_hybrid_swa: + logger.debug(f"Hybrid swa model: {self.hf_config.architectures=}") + + self.is_deepseek_v4_arch = any( + arch + in [ + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", + ] + for arch in self.hf_config.architectures + ) + + if not self.is_deepseek_v4_arch: + self.swa_attention_layer_ids, self.full_attention_layer_ids = ( + get_hybrid_layer_ids( + self.hf_config.architectures, + self.hf_text_config, + ) + ) + + self.has_attention_sinks = self._detect_attention_sinks() + + self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [ + *MIMO_V2_MODEL_ARCHS, + "MiMoV2MTP", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "EmbeddingGemma2Model", + "InklingForConditionalGeneration", + "InklingForConditionalGenerationMTP", + "Gemma4UnifiedForConditionalGeneration", + "DiffusionGemmaForBlockDiffusion", + ] + + @cached_property + def linear_attn_registry_result(self) -> Any: + return get_linear_attn_config(self.hf_config) + + @property + def has_asymmetric_kv(self) -> bool: + """Whether K and V rows differ in width (MiMoV2 is 192 / 128). + + Not an ``__init__`` field because the MLA special-casing below still + rewrites ``v_head_dim``. + """ + return ( + self.head_dim != self.v_head_dim or self.swa_head_dim != self.swa_v_head_dim + ) + + def _detect_attention_sinks(self) -> bool: + """Check whether the model uses learned attention sinks. + + Attention sinks are per-head scalars added to the softmax denominator + to compensate for evicted KV-cache entries under sliding-window + attention. Not every hybrid-SWA model uses them. + """ + archs = self.hf_config.architectures or [] + if any(a in SWA_SINK_ARCHS for a in archs): + return True + + # MiMoV2 creates sinks only when the config flags are set. + if any(a in archs for a in (*MIMO_V2_MODEL_ARCHS, "MiMoV2MTP")): + return getattr( + self.hf_text_config, "add_swa_attention_sink_bias", False + ) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False) + return bool(getattr(self.hf_text_config, "learnable_sink", False)) + + def _derive_context_length(self, context_length: int): + is_draft_model = self.is_draft_model + derived_context_len = get_context_length(self.hf_text_config) + + if context_length is not None: + if context_length > derived_context_len: + reason = "Target model's" if is_draft_model else "User-specified" + msg = ( + f"Warning: {reason} context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). " + f"This may lead to incorrect model outputs or CUDA errors. Note that the derived context_length may differ from max_position_embeddings in the model's config." + ) + if ( + envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.get() + or is_in_ci() # FIXME: fix this special case + ): + logger.warning(msg) + self.context_len = context_length + if is_draft_model: + self.hf_text_config.max_position_embeddings = context_length + logger.warning( + f"Overriding the draft model's max_position_embeddings to {context_length}." + ) + else: + raise ValueError( + f"{msg} To allow overriding this maximum, set the env var SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1" + ) + else: + self.context_len = context_length + else: + if ( + self.hf_config.architectures + and self.hf_config.architectures[0] == "EmbeddingGemma2Model" + ): + self.context_len = min(derived_context_len, 8192) + else: + self.context_len = derived_context_len + + # Transfer context_len to HuggingFace config so models can access it + self.hf_config.context_len = self.context_len + + def _derive_model_shapes(self): + from sglang.srt.configs.dots3 import Dots3Config + + is_dflash_mla_draft = False + if getattr(self.hf_config, "dflash_config", None): + from sglang.srt.speculative.dflash_utils import ( + parse_dflash_draft_config, + validate_dflash_mla_draft_config, + ) + + is_dflash_mla_draft = parse_dflash_draft_config( + draft_hf_config=self.hf_config + ).is_mla + if is_dflash_mla_draft: + # Validate before the generic defaults below fill missing fields. + validate_dflash_mla_draft_config(self.hf_config) + + # Pure SSM models have no attention heads; use head_dim == 0 so the + # KV-cell size is 0 rather than a division on a missing head count. + is_pure_ssm = any( + arch in self.hf_config.architectures for arch in PURE_SSM_ARCHITECTURES + ) + + # Unify the config keys for hf_text_config + self.head_dim = getattr(self.hf_text_config, "head_dim", None) + if self.head_dim is None: + if is_pure_ssm: + self.head_dim = 0 + else: + self.head_dim = ( + self.hf_text_config.hidden_size + // self.hf_text_config.num_attention_heads + ) + setattr(self.hf_text_config, "head_dim", self.head_dim) + + self.v_head_dim = getattr(self.hf_text_config, "v_head_dim", None) + if self.v_head_dim is None or self.v_head_dim == 0: + self.v_head_dim = self.head_dim + setattr(self.hf_text_config, "v_head_dim", self.v_head_dim) + + self.swa_head_dim = getattr(self.hf_text_config, "swa_head_dim", None) + if self.swa_head_dim is None: + self.swa_head_dim = self.head_dim + setattr(self.hf_text_config, "swa_head_dim", self.swa_head_dim) + + self.swa_v_head_dim = getattr(self.hf_text_config, "swa_v_head_dim", None) + if self.swa_v_head_dim is None: + self.swa_v_head_dim = self.swa_head_dim + setattr(self.hf_text_config, "swa_v_head_dim", self.swa_v_head_dim) + # FIXME: temporary special judge for MLA architecture + if ( + is_dflash_mla_draft + or "DeepseekV2ForCausalLM" in self.hf_config.architectures + or "DeepseekV32ForCausalLM" in self.hf_config.architectures + or "DeepseekV3ForCausalLM" in self.hf_config.architectures + or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures + or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures + or "Glm4MoeLiteForCausalLMNextN" in self.hf_config.architectures + or "GlmMoeDsaForCausalLM" in self.hf_config.architectures + or "GlmMoeDsaForCausalLMNextN" in self.hf_config.architectures + or "Glm5NextForConditionalGeneration" in self.hf_config.architectures + or "Glm5NextForConditionalGenerationNextN" in self.hf_config.architectures + or "LongcatFlashForCausalLM" in self.hf_config.architectures + or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures + or "HYV4ForCausalLM" in self.hf_config.architectures + or "HYV4ForCausalLMNextN" in self.hf_config.architectures + or "DotsVLMForCausalLM" in self.hf_config.architectures + or "Dots3NoteForCausalLM" in self.hf_config.architectures + or "Dots3NoteForCausalLMNextN" in self.hf_config.architectures + or "MistralLarge3ForCausalLM" in self.hf_config.architectures + or ( + "PixtralForConditionalGeneration" in self.hf_config.architectures + and getattr(self.hf_text_config, "kv_lora_rank", None) is not None + ) + or "MistralLarge3ForCausalLMEagle" in self.hf_config.architectures + or "KimiK25ForConditionalGeneration" in self.hf_config.architectures + or "Eagle3DeepseekV2ForCausalLM" in self.hf_config.architectures + or "GigaChat35ForCausalLM" in self.hf_config.architectures + or "GigaChat35ForCausalLMNextN" in self.hf_config.architectures + ): + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + self.v_head_dim = self.hf_text_config.v_head_dim + if isinstance(self.hf_text_config, Dots3Config): + self.swa_kv_lora_rank = self.hf_text_config.swa_kv_lora_rank + self.swa_qk_rope_head_dim = self.hf_text_config.swa_qk_rope_head_dim + else: + self.swa_kv_lora_rank = self.kv_lora_rank + self.swa_qk_rope_head_dim = self.qk_rope_head_dim + if is_dflash_mla_draft: + # Sliding draft layers share the latent rows, not a GQA head_dim. + self.swa_v_head_dim = self.v_head_dim + self.index_head_dim = ( + get_dsa_index_head_dim(self.hf_text_config) + if is_deepseek_dsa(self.hf_text_config) + else None + ) + # In transformers v5, rope_scaling is just rope_parameters. + rope_scaling = getattr( + self.hf_text_config, "rope_parameters", None + ) or getattr(self.hf_text_config, "rope_scaling", None) + self._init_mla_scaling(rope_scaling) + elif ( + "DeepseekV4ForCausalLM" in self.hf_config.architectures + or "DeepseekV4ForCausalLMNextN" in self.hf_config.architectures + or "DeepseekV4ForCausalLMDSpark" in self.hf_config.architectures + ): + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + self.qk_nope_head_dim = self.hf_config.head_dim - self.qk_rope_head_dim + self.window_size = self.hf_config.sliding_window + self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = self.head_dim + self.index_head_dim = self.hf_config.index_head_dim + self.compress_ratios = self.hf_config.compress_ratios + self.attention_arch = AttentionArch.MHA + self._init_mla_scaling(self.hf_config.rope_scaling) + elif "Glm4MoeForCausalLMNextN" in self.hf_config.architectures: + if self.head_dim is None: + self.head_dim = ( + self.hf_text_config.hidden_size + // self.hf_text_config.num_attention_heads + ) + if self.swa_head_dim is None: + self.swa_head_dim = self.head_dim + self.v_head_dim = self.head_dim + self.swa_v_head_dim = self.swa_head_dim + self.attention_arch = AttentionArch.MHA + elif "MiniCPM3ForCausalLM" in self.hf_config.architectures: + self.head_dim = 128 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + elif "DeepseekVL2ForCausalLM" in self.hf_config.architectures and getattr( + self.hf_text_config, "use_mla", True + ): + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim + elif "KimiVLForConditionalGeneration" in self.hf_config.architectures: + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + self.v_head_dim = self.hf_text_config.v_head_dim + self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim + elif ( + "KimiLinearForCausalLM" in self.hf_config.architectures + or "KimiK3LinearForCausalLM" in self.hf_config.architectures + or "KimiK3ForConditionalGeneration" in self.hf_config.architectures + ): + tc = self.hf_text_config + self.head_dim = 72 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = tc.kv_lora_rank + self.qk_rope_head_dim = tc.qk_rope_head_dim + self.v_head_dim = tc.v_head_dim + self.qk_nope_head_dim = tc.qk_nope_head_dim + self._init_mla_scaling(getattr(tc, "rope_scaling", None)) + elif "BailingMoeV2_5ForCausalLM" in self.hf_config.architectures or ( + # Every Bailing draft is renamed to BailingMoeForCausalLMNextN; only + # the MLA ones (V2.5 and V3) carry kv_lora_rank, V2 keeps its GQA + # shapes. + "BailingMoeForCausalLMNextN" in self.hf_config.architectures + and getattr(self.hf_text_config, "kv_lora_rank", None) is not None + ): + self.head_dim = self.hf_text_config.head_dim + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + self.v_head_dim = self.hf_config.v_head_dim + self._init_mla_scaling(self.hf_config.rope_scaling) + elif ( + "BailingMoeV3ForCausalLM" in self.hf_config.architectures + or "BailingMoeV3VLForConditionalGeneration" in self.hf_config.architectures + ): + self.head_dim = 128 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_rope_head_dim = ( + 0 + if getattr(self.hf_text_config, "use_mla_nope", False) + else self.hf_text_config.qk_rope_head_dim + ) + self.v_head_dim = self.hf_text_config.v_head_dim + self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim + self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) + elif ( + "SarvamMLAForCausalLM" in self.hf_config.architectures + or "Glm5NextForConditionalGeneration" in self.hf_config.architectures + ): + self.head_dim = ( + self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim + ) + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim + self.v_head_dim = self.hf_config.v_head_dim + self._init_mla_scaling(self.hf_config.rope_scaling) + else: + if ( + "MistralModel" in self.hf_config.architectures + or "MixtralForCausalLM" in self.hf_config.architectures + or "MistralForCausalLM" in self.hf_config.architectures + ): + if getattr(self, "head_dim", None) is None: + self.head_dim = ( + self.hf_config.hidden_size // self.hf_config.num_attention_heads + ) + # In transformers==4.52.3, the head_dim is null in MistralConfig + if ( + not hasattr(self.hf_text_config, "head_dim") + or self.hf_text_config.head_dim is None + ): + setattr(self.hf_text_config, "head_dim", self.head_dim) + + elif "BaichuanForCausalLM" in self.hf_config.architectures: + self.use_alibi = self.hf_config.hidden_size != 4096 + + # Pure Mamba SSMs have no attention (head_dim set to 0 above). + if is_pure_ssm: + self.attention_arch = AttentionArch.SSM + else: + self.attention_arch = AttentionArch.MHA + + # Mamba2 has no num_attention_heads. + self.num_attention_heads = getattr( + self.hf_text_config, "num_attention_heads", None + ) + self.num_key_value_heads = getattr( + self.hf_text_config, "num_key_value_heads", None + ) + self.first_k_dense_replace = getattr( + self.hf_text_config, "first_k_dense_replace", None + ) + self.full_attention_interval = getattr( + self.hf_text_config, "full_attention_interval", None + ) + + # for Dbrx and MPT models + if self.hf_config.model_type in ["dbrx", "mpt"]: + self.num_key_value_heads = getattr( + self.hf_config.attn_config, "kv_n_heads", None + ) + + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + self.hidden_size = self.hf_text_config.hidden_size + hc_mult = getattr(self.hf_text_config, "hc_mult", 1) + is_glm5_next = getattr(self.hf_config, "model_type", None) == "glm5_next" or ( + getattr(self.hf_text_config, "model_type", None) == "glm5_next_text" + ) + if is_glm5_next and not getattr(self.hf_text_config, "mhc", False): + hc_mult = 1 + if is_glm5_next: + # mHC-flattened hidden size; None when not running an mHC model. + self.hc_hidden_size = self.hidden_size * hc_mult if hc_mult > 1 else None + self.spec_hidden_size = self.hidden_size + else: + self.spec_hidden_size, self.hc_hidden_size = resolve_spec_hidden_size( + self.hf_config, self.hidden_size, hc_mult + ) + self.num_hidden_layers = self.hf_text_config.num_hidden_layers + self.num_attention_layers = self.num_hidden_layers + if "LongcatFlashForCausalLM" in self.hf_config.architectures: + self.num_attention_layers = self.num_hidden_layers * 2 + if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures: + loop_num = getattr(self.hf_text_config, "loop_num", 1) + self.num_attention_layers = int(self.num_hidden_layers * int(loop_num)) + if "NanbeigeForCausalLM" in self.hf_config.architectures: + num_loops = getattr(self.hf_text_config, "num_loops", 1) + self.num_attention_layers = int(self.num_hidden_layers * int(num_loops)) + if "WhisperForConditionalGeneration" in self.hf_config.architectures: + # Whisper has unique layer ID scheme: + # - Encoder self-attention: 0 to encoder_layers-1 (no KV cache) + # - Decoder self-attention: encoder_layers to encoder_layers+decoder_layers-1 (uses KV cache) + # - Decoder cross-attention: encoder_layers+decoder_layers to encoder_layers+2*decoder_layers-1 + # Even though cross-attention doesn't save KV cache, attention backend needs buffer to exist + encoder_layers = getattr(self.hf_text_config, "encoder_layers", 0) + decoder_layers = getattr( + self.hf_text_config, "decoder_layers", self.num_hidden_layers + ) + self.num_attention_layers = encoder_layers + 2 * decoder_layers + self.num_nextn_predict_layers = getattr( + self.hf_text_config, "num_nextn_predict_layers", None + ) + self.vocab_size = self.hf_text_config.vocab_size + # GLM-Image is the only model here whose output head predicts vision tokens. + # Use vision_vocab_size for lm_head, LogitsProcessor, and graph-mode logits buffers. + if _hf_arch(self.hf_config) == "GlmImageForConditionalGeneration": + self.vocab_size = self.hf_text_config.vision_vocab_size + + def _init_mla_scaling(self, rope_scaling: Optional[dict]) -> None: + """Base MLA attention scale from the head dims, then the rope mscale.""" + self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) + if rope_scaling: + self.scaling = compute_mla_mscale_scaling(rope_scaling, self.scaling) + + def get_total_num_attention_heads(self) -> int: + return self.num_attention_heads + + def get_num_attention_heads(self, tensor_parallel_size) -> int: + # Pure-SSM (Mamba) models have no attention; num_attention_heads is None. + total_num_attention_heads = self.num_attention_heads or 0 + return max(1, total_num_attention_heads // tensor_parallel_size) + + # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289 + def get_total_num_kv_heads(self) -> int: + """Returns the total number of KV heads.""" + # For GPTBigCode & Falcon: + # NOTE: for falcon, when new_decoder_architecture is True, the + # multi_query flag is ignored and we use n_head_kv for the number of + # KV heads. + falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"] + new_decoder_arch_falcon = ( + self.hf_config.model_type in falcon_model_types + and getattr(self.hf_config, "new_decoder_architecture", False) + ) + if not new_decoder_arch_falcon and getattr( + self.hf_text_config, "multi_query", False + ): + # Multi-query attention, only one KV head. + # Currently, tensor parallelism is not supported in this case. + return 1 + + # For DBRX and MPT + if self.hf_config.model_type in ["mpt"]: + if "kv_n_heads" in self.hf_config.attn_config: + return self.hf_config.attn_config["kv_n_heads"] + return self.hf_config.num_attention_heads + if self.hf_config.model_type in ["dbrx"]: + return getattr( + self.hf_config.attn_config, + "kv_n_heads", + self.hf_config.num_attention_heads, + ) + if self.hf_config.model_type in ["nemotron-nas"]: + nkvh = { + self.hf_config.num_attention_heads // block.attention.n_heads_in_group + for block in self.hf_config.block_configs + if not block.attention.no_op + } + if len(nkvh) == 0: + raise RuntimeError("Couldn't determine number of kv heads") + if len(nkvh) > 1: + raise ValueError( + "Variable GQA (VGQA) is not yet supported for nemotron-nas in sglang" + ) + return next(iter(nkvh)) + + attributes = [ + # For Falcon: + "n_head_kv", + "num_kv_heads", + # For LLaMA-2: + "num_key_value_heads", + # For ChatGLM: + "multi_query_group_num", + # For Step3 + "num_attention_groups", + ] + for attr in attributes: + num_kv_heads = getattr(self.hf_text_config, attr, None) + if num_kv_heads is not None: + return num_kv_heads + + # Mamba SSMs have no attention, so no KV heads. + if self.attention_arch == AttentionArch.SSM: + return 0 + # For non-grouped-query attention models, the number of KV heads is + # equal to the number of attention heads. + return self.hf_text_config.num_attention_heads + + def get_max_num_attention_heads(self) -> int: + """Max per-layer query head count; num_attention_heads unless the + model sets num_attention_heads_per_layer.""" + per_layer = getattr(self.hf_text_config, "num_attention_heads_per_layer", None) + if per_layer: + return max(per_layer) + return self.num_attention_heads + + def get_num_kv_heads(self, tensor_parallel_size: int, dcp_size: int = 1) -> int: + """Number of KV heads per GPU. + + DCP ranks replicate KV, so heads shard across ``tp // dcp`` groups. + Drafts never join the group and ignore ``dcp_size``. With fewer heads + than groups, each GPU keeps one. + """ + total_num_kv_heads = self.get_total_num_kv_heads() + if self.is_draft_model: + dcp_size = 1 + kv_tensor_parallel_size = tensor_parallel_size // dcp_size + return max(1, total_num_kv_heads // kv_tensor_parallel_size) + + def get_swa_num_kv_heads(self, tensor_parallel_size) -> int: + """Similar to get_num_kv_heads(), but for SWA.""" + if hasattr(self.hf_text_config, "swa_num_key_value_heads"): + total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads + return max(1, total_num_kv_heads // tensor_parallel_size) + elif hasattr(self.hf_text_config, "attention_other_setting"): # For step3p5 + total_num_kv_heads = self.hf_text_config.attention_other_setting.get( + "num_attention_groups" + ) + return max(1, total_num_kv_heads // tensor_parallel_size) + else: + return self.get_num_kv_heads(tensor_parallel_size) + + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py + def _parse_quant_hf_config(self): + quant_cfg = _quant_config_to_dict( + getattr(self.hf_config, "quantization_config", None) + ) + if quant_cfg is not None: + # Identify modelopt quantization + if ( + "quant_method" not in quant_cfg + or quant_cfg["quant_method"] == "modelopt" + ): + parsed_cfg = self._parse_modelopt_quant_config( + {"quantization": quant_cfg} + ) + if parsed_cfg: + quant_cfg.update(parsed_cfg) + + if quant_cfg is None: + # compressed-tensors uses a "compression_config" key + quant_cfg = getattr(self.hf_config, "compression_config", None) + if quant_cfg is None: + # check if is modelopt or mixed-precision model -- Both of them don't have corresponding field + # in hf `config.json` but has a standalone `hf_quant_config.json` in the root directory + # example: https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8/tree/main + # example: https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/tree/main + is_local = os.path.exists(self.model_path) + if not is_local: + # Conditional import based on SGLANG_USE_MODELSCOPE environment variable + if envs.SGLANG_USE_MODELSCOPE.get(): + from modelscope import HubApi, model_file_download + + hf_api = HubApi() + else: + import huggingface_hub + from huggingface_hub import HfApi, hf_hub_download + + hf_api = HfApi() + try: + # In offline mode, skip file_exists check to avoid OfflineModeIsEnabled error + # Instead, directly try to download/read from cache with local_files_only + file_exists = False # Initialize to avoid UnboundLocalError + if not huggingface_hub.constants.HF_HUB_OFFLINE: + # Online mode: check if file exists before attempting download (optimization) + file_exists = retry( + lambda: hf_api.file_exists( + self.model_path, "hf_quant_config.json" + ), + max_retry=2, + initial_delay=1.0, + max_delay=5.0, + ) + if not file_exists: + # File doesn't exist on hub, no need to try downloading + return quant_cfg # None + + # Download (online mode) or read from cache (offline mode) + if envs.SGLANG_USE_MODELSCOPE.get(): + quant_config_file = model_file_download( + model_id=self.model_path, + file_path="hf_quant_config.json", + revision=self.revision, + ) + else: + quant_config_file = hf_hub_download( + repo_id=self.model_path, + filename="hf_quant_config.json", + revision=self.revision, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + with open(quant_config_file) as f: + quant_config_dict = json.load(f) + quant_cfg = self._parse_modelopt_quant_config(quant_config_dict) + except huggingface_hub.errors.LocalEntryNotFoundError: + # Offline mode and file not in cache - this is normal for non-quantized models + logger.debug( + f"hf_quant_config.json not found in cache for {self.model_path} " + "(offline mode, normal for non-quantized models)" + ) + except huggingface_hub.errors.OfflineModeIsEnabled: + # Should not reach here after our changes, but keep for safety + logger.warning( + "Offline mode is enabled, skipping hf_quant_config.json check" + ) + except Exception as e: + logger.warning( + "Failed to load hf_quant_config.json for model %s: %s", + self.model_path, + e, + ) + elif os.path.exists(os.path.join(self.model_path, "hf_quant_config.json")): + quant_config_file = os.path.join( + self.model_path, "hf_quant_config.json" + ) + with open(quant_config_file) as f: + quant_config_dict = json.load(f) + quant_cfg = self._parse_modelopt_quant_config(quant_config_dict) + return quant_cfg + + def _find_quant_modelslim_config(self): + if self.is_draft_model: + return None + quant_config_file = Path(self.model_path, "quant_model_description.json") + quant_cfg = None + if quant_config_file.is_file(): + with open(quant_config_file) as f: + quant_cfg = json.load(f) + # This field is required for flagless model loading but is not present in + # modelslim model description, so we're adding it here manually. + quant_cfg["quant_method"] = "modelslim" + + return quant_cfg + + def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict]: + """Parse ModelOpt quantization config and return the appropriate quant_method. + + Supports both nested LLM ``hf_quant_config.json``:: + + {"quantization": {"quant_algo": "FP8", ...}} + + and flat formats (``config.json`` ``quantization_config``, or diffusion / + unified ModelOpt exports such as Cosmos3):: + + {"quant_algo": "FP8", "quant_method": "modelopt", ...} + """ + json_quant_configs = unwrap_modelopt_quantization_config(quant_config_dict) + if not json_quant_configs: + return None + quant_algo = json_quant_configs.get("quant_algo", None) + + if quant_algo == "MIXED_PRECISION": + quantized_layers = json_quant_configs.get("quantized_layers") or {} + has_modelopt_nvfp4_layers = any( + str(layer_info.get("quant_algo", "")).upper() + in ("NVFP4", "W4A16_NVFP4") + for layer_info in quantized_layers.values() + if isinstance(layer_info, dict) + ) + if has_modelopt_nvfp4_layers: + return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo} + return {"quant_method": "w4afp8", "quant_algo": quant_algo} + elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo): + return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo} + elif quant_algo == "FP8": + return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo} + elif quant_algo == "MXFP8": + group_size = json_quant_configs.get("group_size", 32) + ignored_layers = json_quant_configs.get( + "exclude_modules", json_quant_configs.get("ignore") + ) + kv_cache_quant_algo = json_quant_configs.get("kv_cache_quant_algo") + if kv_cache_quant_algo is None: + kv_cache_scheme = json_quant_configs.get("kv_cache_scheme") + if ( + isinstance(kv_cache_scheme, dict) + and kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + ): + kv_cache_quant_algo = "FP8" + parsed = { + "quant_method": "mxfp8", + "quant_algo": quant_algo, + "activation_scheme": "dynamic", + "weight_block_size": [1, group_size], + "scale_fmt": "ue8m0", + } + if ignored_layers is not None: + parsed["modules_to_not_convert"] = ignored_layers + if kv_cache_quant_algo is not None: + parsed["kv_cache_quant_algo"] = kv_cache_quant_algo + return parsed + else: + return None + + def get_quantization_config_log_str(self) -> Optional[str]: + """ + Get a concise string representation of the quantization config for logging. + Returns something like "quant=fp8, fmt=e4m3" or "quant=gptq, bits=4". + """ + try: + quant_cfg = self._parse_quant_hf_config() + if not quant_cfg: + return None + + quant_method = quant_cfg.get("quant_method", "quantized") + log_str = f"quant={quant_method}" + + # Append interesting fields if they exist + for field in ["bits", "quant_algo", "fmt", "requantization_method"]: + if field in quant_cfg: + log_str += f", {field}={quant_cfg[field]}" + + return log_str + except Exception: + return None + + def _is_already_quantized(self) -> bool: + """Check if the model is already quantized based on config files.""" + # Check for quantization in hf_config (config.json) + if getattr(self.hf_config, "quantization_config", None) or getattr( + self.hf_config, "compression_config", None + ): + return True + + # Check for HuggingFace quantization config + quant_cfg = getattr(self.hf_config, "quantization_config", None) + if quant_cfg is None: + from sglang.srt.utils import has_hf_quant_config + + return has_hf_quant_config(self.model_path) + return True + + def _get_modelopt_quant_type(self) -> str: + """Extract ModelOpt quantization type from unified quantization flag.""" + if self.quantization == "modelopt_fp8": + return "fp8" + elif self.quantization == "modelopt_fp4": + return "nvfp4" + elif self.quantization == "modelopt_mixed": + raise ValueError( + "modelopt_mixed is only supported for pre-quantized checkpoints." + ) + elif self.quantization == "modelopt": + # Auto-detect from model config + quant_cfg = self._parse_quant_hf_config() + if quant_cfg: + quant_method = quant_cfg.get("quant_method", "").lower() + if "fp4" in quant_method: + return "fp4" + elif "fp8" in quant_method: + return "fp8" + # Default to fp8 if can't detect + return "fp8" + else: + return "fp8" # Default fallback + + def _get_sliding_window_size(self) -> Optional[int]: + for key in ("sliding_window_size", "sliding_window", "window_size"): + value = getattr(self.hf_text_config, key, None) + if value is not None: + return value + return None + + def _validate_quantize_and_serve_config(self): + """Validate quantize_and_serve configuration.""" + if not self.quantize_and_serve: + return + + # Check if ModelOpt quantization is specified + _MODELOPT_QUANTIZATION_METHODS = [ + "modelopt", + "modelopt_fp8", + "modelopt_fp4", + "nvfp4_online", + "modelopt_mixed", + ] + modelopt_quantization_specified = ( + self.quantization in _MODELOPT_QUANTIZATION_METHODS + ) + + if not modelopt_quantization_specified: + raise ValueError( + "quantize_and_serve requires ModelOpt quantization (set with --quantization " + f"{{{', '.join(sorted(_MODELOPT_QUANTIZATION_METHODS))}}})" + ) + + # quantize_and_serve is disabled due to compatibility issues + raise NotImplementedError( + "quantize_and_serve functionality is currently disabled due to compatibility issues. " + "Please use the separate quantize-then-deploy workflow instead. " + "Step 1: Quantize and export model. " + "Step 2: Deploy the exported model." + ) + + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py + def _verify_quantization(self) -> None: + supported_quantization = [*QUANTIZATION_METHODS] + rocm_supported_quantization = [ + "awq", + "fp8", + "compressed_tensors", + "compressed-tensors", + "w8a8_fp8", + "petit_nvfp4", + "quark", + "mxfp4", + "mxfp8", + "auto-round", + "auto-round-int8", + "quark_int4fp8_moe", + "quark_mxfp4", + ] + optimized_quantization_methods = [ + "fp8", + "marlin", + "modelopt_fp8", + "modelopt_fp4", + "modelopt_mixed", + "nvfp4_online", + "gptq_marlin_24", + "gptq_marlin", + "awq_marlin", + "compressed_tensors", + "compressed-tensors", + "experts_int8", + "w8a8_int8", + "w8a8_fp8", + "moe_wna16", + "w4afp8", + "petit_nvfp4", + "quark", + "modelslim", + "humming", + "quark_mxfp4", + "auto-round", + "minimax_m3_nvfp4", + ] + compatible_quantization_methods = { + "modelopt_fp8": ["modelopt"], + # Keep explicit or inherited modelopt_fp4 for literal FP8 checkpoints + # so eligible MoE experts are requantized online. + "modelopt_fp4": ["modelopt", "fp8"], + "modelopt_mixed": ["modelopt"], + "nvfp4_online": ["fp8", "modelopt_fp8"], + # Resolved by the MiniMax override from the MoE fields. + "minimax_m3_nvfp4": ["mxfp8"], + "petit_nvfp4": ["modelopt"], + "w8a8_int8": ["compressed-tensors", "compressed_tensors"], + "w8a8_fp8": ["compressed-tensors", "compressed_tensors"], + "auto-round-int8": ["compressed-tensors", "compressed_tensors"], + } + if self.quantization is not None: + self.quantization = self.quantization.lower() + + # Parse quantization method from the HF and ModelSlim model config, if available. + # Only one function should return config, other should return None. + cfg_list = [] + hf_config = self._parse_quant_hf_config() + modelslim_config = self._find_quant_modelslim_config() + quant_config = modelslim_config or hf_config + if quant_config is not None: + cfg_list.append(quant_config) + + # Filter out None values + cfg_list = [item for item in cfg_list if item is not None] + if len(cfg_list) > 1: + raise ValueError( + "Config list contains configs from 2 methods, must be only 1" + ) + quant_cfg = cfg_list[0] if cfg_list else None + + if quant_cfg is not None: + quant_method = quant_cfg.get( + "quant_method", "" if not self.quantization else self.quantization + ).lower() + + # ModelOpt FP4 and mixed checkpoints can quantize only the target + # model; an embedded MTP draft may stay unquantized, so an explicit + # nvfp4_online opt-in for the draft wins over checkpoint detection. + # The online loader rejects already-packed weights at load time. + preserve_online_draft_quantization = ( + self.is_draft_model + and self.quantization == "nvfp4_online" + and quant_method in ("modelopt_fp4", "modelopt_mixed") + ) + # An explicit online-requantization request (e.g. quark_mxfp4 on top + # of an NVFP4/mixed checkpoint) must not be overridden back to the + # source format + if self.quantization not in REQUANTIZATION_METHODS: + # Detect which checkpoint is it + if not preserve_online_draft_quantization: + for _, method in QUANTIZATION_METHODS.items(): + quantization_override = method.override_quantization_method( + quant_cfg, self.quantization + ) + if quantization_override: + quant_method = quantization_override + self.quantization = quantization_override + break + + # Verify quantization configurations. + if self.quantization is None: + self.quantization = quant_method + elif self.quantization != quant_method: + # Check if the CLI-specified quantization is compatible with HF config's quant_method + is_compatible = preserve_online_draft_quantization or ( + self.quantization in compatible_quantization_methods + and quant_method + in compatible_quantization_methods[self.quantization] + ) + if is_compatible: + # Keep the CLI-specified quantization (e.g., modelopt_fp4) even if + # HF config says "modelopt" - they are compatible + logger.info( + f"Using CLI-specified quantization ({self.quantization}) which is " + f"compatible with HF config quant_method ({quant_method})." + ) + elif self.is_draft_model and not ( + self.is_draft_quantization_explicit + and self.quantization in REQUANTIZATION_METHODS + and is_hip() + and is_qwen3_5_mtp_draft(self.hf_config) + ): + # Allow auto-detection of quantization from checkpoint for draft model + # only if the CLI quantization is not compatible. An explicit + # online-requantization request for the draft (e.g. quark_mxfp4 + # for an MTP stack the checkpoint left in bf16) is honored below. + logger.info( + f"Draft model quantization ({quant_method}) differs from " + f"main model quantization ({self.quantization}). " + f"Using draft model's detected quantization: {quant_method}" + ) + self.quantization = quant_method + elif self.quantization in REQUANTIZATION_METHODS: + logger.info_once( + f"Requantizing from quant_method='{quant_method}' to the requested online quantization='{self.quantization}'. Beware that requantization may incur a loss in accuracy, the requantized model should be re-validated/re-evaluated. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization." + ) + else: + raise ValueError( + "Quantization method specified in the model config " + f"({quant_method}) does not match the quantization " + f"method specified in the `quantization` argument " + f"({self.quantization})." + ) + + if self.quantization is not None: + if self.quantization not in supported_quantization: + raise ValueError( + f"Unknown quantization method: {self.quantization}. Must " + f"be one of {supported_quantization}." + ) + if is_hip() and self.quantization not in rocm_supported_quantization: + raise ValueError( + f"{self.quantization} quantization is currently not " + f"supported in ROCm." + ) + if self.quantization not in optimized_quantization_methods: + # Don't warn for MXFP4/MXFP8 on SM100 since they have optimized kernels + if not ( + self.quantization in ["mxfp4", "mxfp8"] and get_platform().is_sm100 + ): + logger.warning( + "%s quantization is not fully " + "optimized yet. The speed can be slower than " + "non-quantized models.", + self.quantization, + ) + + def _verify_dual_chunk_attention_config(self) -> None: + if hasattr(self.hf_config, "dual_chunk_attention_config"): + # Try loading the sparse attention config + sparse_attn_config = get_sparse_attention_config(self.model_path) + if not sparse_attn_config: + return + self.hf_config.dual_chunk_attention_config["sparse_attention_config"] = ( + sparse_attn_config + ) + if ( + "sparse_attention_enabled" + not in self.hf_config.dual_chunk_attention_config + ): + self.hf_config.dual_chunk_attention_config[ + "sparse_attention_enabled" + ] = True + + def _get_hf_eos_token_id(self) -> Optional[Set[int]]: + eos_ids = getattr(self.hf_config, "eos_token_id", None) + if eos_ids is not None: + # it can be either int or list of int + eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) + if eos_ids is None: + eos_ids = set() + if self.hf_generation_config: + generation_eos_ids = getattr( + self.hf_generation_config, "eos_token_id", None + ) + if generation_eos_ids: + generation_eos_ids = ( + {generation_eos_ids} + if isinstance(generation_eos_ids, int) + else set(generation_eos_ids) + ) + eos_ids = eos_ids | generation_eos_ids + return eos_ids + + def get_default_sampling_params(self) -> dict[str, Any]: + """ + Get default sampling parameters from the model's generation config. + + This method returns non-default sampling parameters from the model's + generation_config.json when sampling_defaults is set to "model". + + Returns: + A dictionary containing the non-default sampling parameters. + """ + if self.sampling_defaults != "model": + return {} + + if self.hf_generation_config is None: + return {} + + config = self.hf_generation_config.to_dict() + + available_params = [ + "repetition_penalty", + "temperature", + "top_k", + "top_p", + "min_p", + ] + + default_sampling_params = { + p: config.get(p) for p in available_params if config.get(p) is not None + } + + return default_sampling_params + + def _maybe_pull_model_for_runai(self, model: str) -> None: + if is_runai_obj_uri(model): + # local path for loading the config + self.model_path = ObjectStorageModel.get_path(model) + # remote path for loading the weights + self.model_weights = model + + def _maybe_pull_model_tokenizer_from_remote(self) -> None: + """ + Pull the model config files to a temporary + directory in case of remote. + + Args: + model: The model name or path. + + """ + from sglang.srt.connector import create_remote_connector + from sglang.srt.utils import is_remote_url + + if is_remote_url(self.model_path): + logger.info("Pulling model configs from remote...") + # BaseConnector implements __del__() to clean up the local dir. + # Since config files need to exist all the time, so we DO NOT use + # with statement to avoid closing the client. + client = create_remote_connector(self.model_path) + if is_remote_url(self.model_path): + client.pull_files(allow_pattern=["*config.json"]) + self.model_weights = self.model_path + self.model_path = client.get_local_dir() + + +# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +_STR_DTYPE_TO_TORCH_DTYPE = { + "half": torch.float16, + "float16": torch.float16, + "float": torch.float32, + "float32": torch.float32, + "bfloat16": torch.bfloat16, +} + + +# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +def _get_and_verify_dtype( + config: PretrainedConfig, + dtype: Union[str, torch.dtype], +) -> torch.dtype: + # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct + # because config.torch_dtype can be None. + if isinstance(config, dict): + config_dtype = config.get("dtype", None) or config.get("torch_dtype", None) + model_type = config.get("model_type", "") + else: + config_dtype = getattr(config, "dtype", None) + model_type = getattr(config, "model_type", "") + if isinstance(config_dtype, str): + config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None) + if config_dtype is None: + config_dtype = torch.float32 + + if isinstance(dtype, str): + dtype = dtype.lower() + if dtype == "auto": + if config_dtype == torch.float32: + if model_type.startswith("gemma"): + if model_type == "gemma": + gemma_version = "" + else: + gemma_version = model_type[5] + logger.info( + f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead " + "of float16 by default. Please specify `dtype` if you " + "want to use float16." + ) + torch_dtype = torch.bfloat16 + else: + # Following the common practice, we use float16 for float32 + # models. + torch_dtype = torch.float16 + else: + torch_dtype = config_dtype + else: + if dtype not in _STR_DTYPE_TO_TORCH_DTYPE: + raise ValueError(f"Unknown dtype: {dtype}") + torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype] + elif isinstance(dtype, torch.dtype): + torch_dtype = dtype + else: + raise ValueError(f"Unknown dtype: {dtype}") + + # Verify the dtype. + if torch_dtype != config_dtype: + if torch_dtype == torch.float32: + # Upcasting to float32 is allowed. + logger.debug("Upcasting %s to %s.", config_dtype, torch_dtype) + pass + elif config_dtype == torch.float32: + # Downcasting from float32 to float16 or bfloat16 is allowed. + logger.debug("Downcasting %s to %s.", config_dtype, torch_dtype) + pass + else: + # Casting between float16 and bfloat16 is allowed with a warning. + logger.debug("Casting %s to %s.", config_dtype, torch_dtype) + + return torch_dtype + + +def is_embedding_gemma(config) -> bool: + """Whether ``config`` is Google's bidirectional EmbeddingGemma checkpoint. + + EmbeddingGemma uses the otherwise generative ``Gemma3TextModel`` + architecture, so its model type alone is insufficient for dispatch. The + upstream ``use_bidirectional_attention`` flag is the defining distinction. + """ + return getattr(config, "model_type", None) == "gemma3_text" and getattr( + config, "use_bidirectional_attention", False + ) + + +def is_generation_model(model_architectures: List[str], is_embedding: bool = False): + # We have two ways to determine whether a model is a generative model. + # 1. Check the model architecture + # 2. check the `is_embedding` server args + + if ( + "LlamaEmbeddingModel" in model_architectures + or "MistralModel" in model_architectures + or "LlamaForSequenceClassification" in model_architectures + or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures + or "InternLM2ForRewardModel" in model_architectures + or "Qwen2ForRewardModel" in model_architectures + or "Qwen3ForRewardModel" in model_architectures + or "Qwen2ForSequenceClassification" in model_architectures + or "Qwen3ForSequenceClassification" in model_architectures + or "Qwen3Model" in model_architectures + or "CLIPModel" in model_architectures + or "BertModel" in model_architectures + or "Contriever" in model_architectures + or "BertForSequenceClassification" in model_architectures + or "XLMRobertaModel" in model_architectures + or "XLMRobertaForSequenceClassification" in model_architectures + or "Gemma2ForSequenceClassification" in model_architectures + or "EmbeddingGemma2Model" in model_architectures + or "Lfm2BidirectionalModel" in model_architectures + ): + return False + else: + return not is_embedding + + +multimodal_model_archs = [ + "CLIPModel", + "Cohere2VisionForConditionalGeneration", + "Cosmos3EdgeForConditionalGeneration", + "DeepseekVL2ForCausalLM", + "Ernie4_5_VLMoeForConditionalGeneration", + "MiniMaxM3SparseForConditionalGeneration", + "Gemma3ForConditionalGeneration", + "Gemma3nForConditionalGeneration", + "Gemma4ForConditionalGeneration", + "EmbeddingGemma2Model", + "Gemma4UnifiedForConditionalGeneration", + "DiffusionGemmaForBlockDiffusion", + "Glm4vForConditionalGeneration", + "Glm4vMoeForConditionalGeneration", + "Glm5NextForConditionalGeneration", + "GlmOcrForConditionalGeneration", + "GlmAsrForConditionalGeneration", + "GlmImageForConditionalGeneration", + "Grok1VForCausalLM", + "Grok1AForCausalLM", + "LlavaLlamaForCausalLM", + "Llama4ForConditionalGeneration", + "LlavaMistralForCausalLM", + "LlavaQwenForCausalLM", + "LlavaForConditionalGeneration", + "LlavaVidForCausalLM", + "Lfm2VlForConditionalGeneration", + "LightOnOCRForConditionalGeneration", + *MIMO_V2_MULTIMODAL_ARCHS, + "MiMoV2ASRForCausalLM", + "MiniCPMO", + "MiniCPMV", + "Mistral3ForConditionalGeneration", + "MultiModalityCausalLM", + "MllamaForConditionalGeneration", + "MossVLForConditionalGeneration", + "NemotronH_Nano_VL_V2", + "NemotronH_Nano_Omni_Reasoning_V3", + "NemotronH_Omni_Reasoning_V3", + "MuseGlimmerForConditionalGeneration", + "MuseSparkForConditionalGeneration", + "PixtralForConditionalGeneration", + "Qwen2AudioForConditionalGeneration", + "Qwen2VLForConditionalGeneration", + "Qwen2_5_VLForConditionalGeneration", + "Qwen3VLForConditionalGeneration", + "Qwen3VLMoeForConditionalGeneration", + "Qwen3_5ForConditionalGeneration", + "SenseNovaLoopForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen4ExpForConditionalGeneration", + "InternS2PreviewForConditionalGeneration", + "InternS2MobiusForConditionalGeneration", + "Qwen3ASRForConditionalGeneration", + "Qwen3OmniMoeForConditionalGeneration", + "KimiVLForConditionalGeneration", + "LocateAnythingForConditionalGeneration", + "InternVLChatModel", + "InternS1ForConditionalGeneration", + "InternS1ProForConditionalGeneration", + "Phi4MMForCausalLM", + "VoxtralForConditionalGeneration", + "WhisperForConditionalGeneration", + "Step3VLForConditionalGeneration", + "POINTSV15ChatModel", + "DotsVLMForCausalLM", + "Dots3NoteForCausalLM", + "DotsOCRForCausalLM", + "Sarashina2VisionForCausalLM", + "NVILAForConditionalGeneration", + "NVILALiteForConditionalGeneration", + "DeepseekOCRForCausalLM", + "UnlimitedOCRForCausalLM", + "JetVLMForConditionalGeneration", + "PaddleOCRVLForConditionalGeneration", + "MiDashengLMModel", + "StepVLForConditionalGeneration", + "Step3p7ForConditionalGeneration", + "KimiK25ForConditionalGeneration", + "BailingMMNativeForConditionalGeneration", + "BailingMM2NativeForConditionalGeneration", + "BailingMoeV3VLForConditionalGeneration", +] + +piecewise_cuda_graph_disabled_model_archs = [ + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", + "Qwen3NextForCausalLM", + "Glm5NextForConditionalGeneration", + "BailingMoeV2_5ForCausalLM", + "LLaDAModelLM", +] + +# Multimodal archs allowed to keep prefill piecewise CUDA graph enabled. The +# generic "multimodal model" rule in ServerArgs disables prefill piecewise CG for +# all multimodal models; archs here opt back in because their LM prefill captures +# cleanly (vision encoder runs eagerly outside the graph via general_mm_embed_routine). +multimodal_piecewise_cuda_graph_supported_model_archs = [ + "KimiK25ForConditionalGeneration", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", +] + +# Multimodal archs whose LM prefill is validated under breakable CUDA graph; +# replay eligibility for embed-carrying batches is checked by can_run_graph. +# The Kimi archs are structurally multimodal -- their configs always carry a +# vision_config, so is_multimodal is True even for text-only serving -- and the +# generic multimodal rule disabled prefill CG for them despite the LM prefill +# capturing cleanly. +multimodal_breakable_cuda_graph_supported_model_archs = [ + "MiniMaxM3SparseForConditionalGeneration", + "Cohere2VisionForConditionalGeneration", + "Glm5NextForConditionalGeneration", + "InternS2MobiusForConditionalGeneration", + "PaddleOCRVLForConditionalGeneration", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen4ExpForConditionalGeneration", + "MuseGlimmerForConditionalGeneration", + "MuseSparkForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "KimiK25ForConditionalGeneration", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", +] + +if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): + multimodal_model_archs.append(external_mm_model_arch) + + +def is_multimodal_model(model_architectures: List[str]): + if any( + multi_model_arch in model_architectures + for multi_model_arch in multimodal_model_archs + ): + return True + else: + return False + + +def is_audio_model(model_architectures: List[str]): + models = [ + "WhisperForConditionalGeneration", + "Qwen3ASRForConditionalGeneration", + ] + return any(model in model_architectures for model in models) + + +def is_encoder_decoder_model(model_architectures: List[str]): + models = [ + "WhisperForConditionalGeneration", + "MllamaForConditionalGeneration", + "MossVLForConditionalGeneration", + ] + return any(model in model_architectures for model in models) + + +def is_local_attention_model(model_architectures: List[str]): + return "Llama4ForConditionalGeneration" in model_architectures + + +def is_multimodal_chunked_prefill_supported(model_architectures: List[str]): + """Check if chunked prefill is supported for a MultiModal model.""" + unsupported = [ + "Grok1VForCausalLM", + "Grok1AForCausalLM", + "LlavaLlamaForCausalLM", + "MllamaForConditionalGeneration", + "MossVLForConditionalGeneration", + "CLIPModel", + ] + if any(multi_model_arch in unsupported for multi_model_arch in model_architectures): + return False + else: + return True + + +def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]): + return any( + arch in piecewise_cuda_graph_disabled_model_archs + for arch in model_architectures + ) + + +def is_multimodal_piecewise_cuda_graph_supported(model_architectures: List[str]): + """Whether a multimodal arch may keep prefill piecewise CUDA graph enabled.""" + return any( + arch in multimodal_piecewise_cuda_graph_supported_model_archs + for arch in model_architectures + ) + + +def is_multimodal_breakable_cuda_graph_supported(model_architectures: List[str]): + """Whether a multimodal arch may keep prefill breakable CUDA graph enabled.""" + return any( + arch in multimodal_breakable_cuda_graph_supported_model_archs + for arch in model_architectures + ) + + +# SequenceClassification models that use CrossEncodingPooler +_cross_encoding_pooler_archs = [ + "BertForSequenceClassification", + "XLMRobertaForSequenceClassification", +] + + +def is_cross_encoding_pooler_model(model_architectures: List[str]) -> bool: + return any(arch in _cross_encoding_pooler_archs for arch in model_architectures) + + +# SequenceClassification models whose forward routes the head through +# score_and_pool (per-position pooling); only these support setwise readout +# (token_indices_to_pool). Keep in sync with callers of layers.pooler.score_and_pool. +_score_and_pool_archs = [ + "LlamaForSequenceClassification", + "Qwen2ForSequenceClassification", + "Qwen3ForSequenceClassification", +] + + +def is_score_and_pool_model(model_architectures: List[str]) -> bool: + return any(arch in _score_and_pool_archs for arch in model_architectures) + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float: + """Compute MLA attention scaling factor from rope_scaling with mscale. + + Used by DeepSeek, BailingMoe, SarvamMLA and similar MLA models. + Transformers v5 also exposes the default RoPE parameters through + ``rope_scaling``. Those parameters do not request any scaling. + Warns if 'factor' is missing from a scaling request (common in v5 configs). + """ + # v5 uses "rope_type", v4 uses "type" + rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type") + if rope_type == "default": + return base_scaling + + if not rope_scaling.get("apply_yarn_scaling", True) or not rope_scaling.get( + "apply_scale", True + ): + return base_scaling + + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) + if "factor" not in rope_scaling: + logger.warning( + "rope_scaling missing 'factor', defaulting to 1.0. Check model accuracy.", + ) + scaling_factor = rope_scaling.get("factor", 1.0) + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + return base_scaling * mscale * mscale + + +def is_hybrid_swa_model( + model_architectures: List[str], + hf_text_config: Optional[PretrainedConfig] = None, +): + hybrid_swa_archs = { + "BeamForCausalLM", + "Llama4ForConditionalGeneration", + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLMNextN", + "DeepseekV4ForCausalLMDSpark", + *SWA_SINK_ARCHS, + *MIMO_V2_MODEL_ARCHS, + "MiMoV2MTP", + "Step3p5ForCausalLM", + "Step3p5MTP", + "Step3p7ForConditionalGeneration", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "EmbeddingGemma2Model", + "Gemma4UnifiedForConditionalGeneration", + "DiffusionGemmaForBlockDiffusion", + "LagunaForCausalLM", + "MellumForCausalLM", + "MuseGlimmerForCausalLM", + "MuseGlimmerForConditionalGeneration", + "InklingForConditionalGeneration", + "InklingForConditionalGenerationMTP", + "UnlimitedOCRForCausalLM", + } + if any(arch in hybrid_swa_archs for arch in model_architectures): + if "BeamForCausalLM" in model_architectures and hf_text_config is not None: + # A schedule without sliding layers uses the ordinary KV pool; the + # SWA pool also serves all-sliding schedules. + return "sliding_attention" in hf_text_config.layer_types + # Only treat Laguna as hybrid SWA when it actually has a sliding window. + if ( + "LagunaForCausalLM" in model_architectures + and hf_text_config is not None + and not getattr(hf_text_config, "sliding_window", 0) + ): + return False + return True + # Also recognize models that explicitly opt-in via their HF text config, + # so custom hybrid-SWA architectures don't need to be added to the allowlist. + if hf_text_config is not None and getattr(hf_text_config, "is_hybrid_swa", False): + return True + return False + + +def get_hybrid_layer_ids( + model_architectures: List[str], + hf_text_config: PretrainedConfig, +): + num_hidden_layers = hf_text_config.num_hidden_layers + if "Llama4ForConditionalGeneration" in model_architectures: + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if (i + 1) % 4 != 0 + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if (i + 1) % 4 == 0 + ] + elif any(arch in SWA_SINK_ARCHS for arch in model_architectures) or any( + arch in ("Dots3NoteForCausalLM", "Dots3NoteForCausalLMNextN") + for arch in model_architectures + ): + layer_types = getattr(hf_text_config, "layer_types", []) + swa_attention_layer_ids = [ + i for i, x in enumerate(layer_types) if x == "sliding_attention" + ] + full_attention_layer_ids = [ + i for i, x in enumerate(layer_types) if x == "full_attention" + ] + elif any(arch in MIMO_V2_MODEL_ARCHS for arch in model_architectures): + hybrid_layer_pattern = getattr(hf_text_config, "hybrid_layer_pattern", None) + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1 + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0 + ] + elif "MiMoV2MTP" in model_architectures: + swa_attention_layer_ids = [0] + full_attention_layer_ids = [] + elif ( + "Step3p5ForCausalLM" in model_architectures + or "Step3p7ForConditionalGeneration" in model_architectures + ): + layer_types = hf_text_config.layer_types + swa_attention_layer_ids = [ + i + for i, x in enumerate(layer_types) + if x == "sliding_attention" and i < num_hidden_layers + ] + full_attention_layer_ids = [ + i + for i, x in enumerate(layer_types) + if x == "full_attention" and i < num_hidden_layers + ] + elif "Step3p5MTP" in model_architectures: + swa_attention_layer_ids = [0] + full_attention_layer_ids = [] + elif ( + "BeamForCausalLM" in model_architectures + or "Gemma4ForCausalLM" in model_architectures + or "Gemma4ForConditionalGeneration" in model_architectures + or "EmbeddingGemma2Model" in model_architectures + or "Gemma4UnifiedForConditionalGeneration" in model_architectures + or "DiffusionGemmaForBlockDiffusion" in model_architectures + or "LagunaForCausalLM" in model_architectures + or "MellumForCausalLM" in model_architectures + or "MuseGlimmerForCausalLM" in model_architectures + or "MuseGlimmerForConditionalGeneration" in model_architectures + ): + layer_types = getattr(hf_text_config, "layer_types", []) + swa_attention_layer_ids = [ + i for i, x in enumerate(layer_types) if x == "sliding_attention" + ] + full_attention_layer_ids = [ + i for i, x in enumerate(layer_types) if x == "full_attention" + ] + elif "InklingForConditionalGenerationMTP" in model_architectures: + # Runner-local layer setup selects one depth from these shared lists, + # including when every draft depth uses full attention. + local_set = set(hf_text_config.mtp_local_layer_ids) + swa_attention_layer_ids = sorted(local_set) + full_attention_layer_ids = [ + i + for i in range(hf_text_config.num_nextn_predict_layers) + if i not in local_set + ] + elif "InklingForConditionalGeneration" in model_architectures: + local_layer_ids = hf_text_config.local_layer_ids + local_layer_id_set = set(local_layer_ids) + assert len(local_layer_id_set) == len(local_layer_ids), ( + f"Inkling local_layer_ids must be unique: {local_layer_ids}" + ) + assert all( + 0 <= layer_id < num_hidden_layers for layer_id in local_layer_id_set + ), ( + f"Inkling local_layer_ids must be in [0, {num_hidden_layers}): {local_layer_ids}" + ) + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if i in local_layer_id_set + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if i not in local_layer_id_set + ] + elif "UnlimitedOCRForCausalLM" in model_architectures: + swa_attention_layer_ids = list(range(num_hidden_layers)) + full_attention_layer_ids = [] + elif "MuseSparkForConditionalGeneration" in model_architectures: + swa_attention_layer_ids = [ + i + for i, kind in enumerate(hf_text_config.layer_types) + if kind == "sliding_attention" + ] + full_attention_layer_ids = [ + i + for i, kind in enumerate(hf_text_config.layer_types) + if kind == "full_attention" + ] + elif getattr(hf_text_config, "hybrid_layer_pattern", None) is not None: + # Generic fallback for custom hybrid SWA models that opt in via + # hf_text_config.is_hybrid_swa and expose a hybrid_layer_pattern + # (1 = SWA, 0 = full) without needing to be added to the allowlist. + hybrid_layer_pattern = hf_text_config.hybrid_layer_pattern + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1 + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0 + ] + else: + swa_attention_layer_ids = None + full_attention_layer_ids = None + return swa_attention_layer_ids, full_attention_layer_ids