# 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_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 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) == "DSparkDraftModel"


# 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)


def is_qwen3_5(config) -> bool:
    return _hf_arch(config) in (
        "Qwen3_5ForConditionalGeneration",
        "Qwen3_5MoeForConditionalGeneration",
        "Qwen3_5ForCausalLM",
        "Qwen3_5MoeForCausalLM",
    )


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 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.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)
        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.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.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] 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

        # 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 (
            "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
            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",
        ]
        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"],
            "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",
    "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 = [
    "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 = {
        "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):
        # 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 (
        "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