lightseekorg--tokenspeed/python/tokenspeed/runtime/configs/model_config.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Model configuration helpers and derived runtime metadata."""

import copy
import json
import math
import os
from collections.abc import Callable
from dataclasses import dataclass
from enum import IntEnum, auto

import torch
import yaml
from transformers import PretrainedConfig

from tokenspeed.runtime.layers.quantization import QUANTIZATION_METHODS
from tokenspeed.runtime.utils import get_colorful_logger
from tokenspeed.runtime.utils.env import envs
from tokenspeed.runtime.utils.hf_transformers_utils import (
    get_config,
    get_context_length,
    get_generation_config,
    resolve_architecture,
)
from tokenspeed.runtime.utils.server_args import ServerArgs

logger = get_colorful_logger(__name__)

_DEEPSEEK_V4_ARCHITECTURES = frozenset(
    {
        "DeepseekV4ForCausalLM",
        "DeepseekV4ForCausalLMNextN",
    }
)
_MLA_ARCHITECTURES = frozenset(
    {
        "DeepseekV3ForCausalLM",
        "DeepseekV3ForCausalLMNextN",
        "Eagle3DeepseekV2ForCausalLM",
        "LongcatFlashForCausalLM",
        "KimiK25ForConditionalGeneration",
    }
)
_DSA_ARCHITECTURES = frozenset(
    {
        "GlmMoeDsaForCausalLM",
        "GlmMoeDsaForCausalLMNextN",
    }
)
_DOUBLE_ATTENTION_LAYER_ARCHITECTURES = frozenset(
    {
        "LongcatFlashForCausalLM",
    }
)


class AttentionArch(IntEnum):
    MLA = auto()
    MHA = auto()
    DSA = auto()


@dataclass(frozen=True)
class _AttentionFamilySpec:
    name: str
    architectures: frozenset[str]
    configure: Callable[[object], None]
    default_backend: str | None = None
    supports_target_verify_forward_mode: bool = False
    default_block_size: int | None = None


def override_model_config(model_config, ext_yaml):
    with open(ext_yaml, encoding="utf-8") as f:
        ext_config = yaml.safe_load(f)

    override_model_config: dict = ext_config.get("override_model_config", {})
    for k, v in override_model_config.items():
        if hasattr(model_config, k):
            old_v = model_config.__getattribute__(k)
            if isinstance(v, dict):
                new_v = copy.deepcopy(old_v)
                new_v.__dict__.update(v)
            else:
                new_v = v
            model_config.__setattr__(k, new_v)
            logger.info("Override model config: %s=%r", k, new_v)


def is_deepseek_v4(config: PretrainedConfig) -> bool:
    return resolve_architecture(config) in _DEEPSEEK_V4_ARCHITECTURES


def is_deepseek_v4_nextn(config: PretrainedConfig) -> bool:
    return resolve_architecture(config) == "DeepseekV4ForCausalLMNextN"


def configure_deepseek_v4_attention(model_config) -> None:
    """Derive DeepSeek V4's MLA-like dimensions for runtime setup."""

    hf_config = model_config.hf_config
    model_config.head_dim = hf_config.head_dim
    model_config.attention_arch = AttentionArch.MLA
    model_config.kv_lora_rank = hf_config.head_dim
    model_config.qk_rope_head_dim = hf_config.qk_rope_head_dim
    model_config.qk_nope_head_dim = hf_config.head_dim - hf_config.qk_rope_head_dim
    model_config.v_head_dim = hf_config.head_dim
    model_config.index_head_dim = getattr(hf_config, "index_head_dim", None)
    model_config.scaling = 1 / math.sqrt(model_config.head_dim)
    rope_scaling = getattr(hf_config, "rope_scaling", None)
    if rope_scaling:
        mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
        scaling_factor = rope_scaling["factor"]
        mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
        model_config.scaling = model_config.scaling * mscale * mscale


def configure_glm_attention(model_config) -> None:
    mla_config = (
        model_config.hf_text_config
        if hasattr(model_config.hf_text_config, "kv_lora_rank")
        else model_config.hf_config
    )
    required_fields = (
        "kv_lora_rank",
        "qk_nope_head_dim",
        "qk_rope_head_dim",
        "v_head_dim",
        "index_topk",
        "index_head_dim",
        "index_n_heads",
    )
    missing_fields = [
        field for field in required_fields if not hasattr(mla_config, field)
    ]
    if missing_fields:
        raise ValueError(
            "GLM attention config is missing required fields: "
            + ", ".join(missing_fields)
        )

    model_config.head_dim = getattr(mla_config, "qk_head_dim", None)
    if model_config.head_dim is None:
        model_config.head_dim = (
            mla_config.qk_nope_head_dim + mla_config.qk_rope_head_dim
        )
    model_config.attention_arch = AttentionArch.DSA
    model_config.kv_lora_rank = mla_config.kv_lora_rank
    model_config.qk_nope_head_dim = mla_config.qk_nope_head_dim
    model_config.qk_rope_head_dim = mla_config.qk_rope_head_dim
    model_config.v_head_dim = mla_config.v_head_dim
    model_config.index_topk = mla_config.index_topk
    model_config.index_head_dim = mla_config.index_head_dim
    model_config.index_n_heads = mla_config.index_n_heads
    model_config.index_topk_pattern = getattr(mla_config, "index_topk_pattern", None)

    model_config.scaling = 1 / math.sqrt(
        model_config.qk_nope_head_dim + model_config.qk_rope_head_dim
    )
    rope_scaling = getattr(mla_config, "rope_scaling", None)
    if rope_scaling and "factor" in rope_scaling:
        mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
        scaling_factor = rope_scaling["factor"]
        mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
        model_config.scaling = model_config.scaling * mscale * mscale


def configure_mla_attention(model_config) -> None:
    mla_config = (
        model_config.hf_text_config
        if hasattr(model_config.hf_text_config, "kv_lora_rank")
        else model_config.hf_config
    )
    model_config.head_dim = 256
    model_config.attention_arch = AttentionArch.MLA
    model_config.kv_lora_rank = mla_config.kv_lora_rank
    model_config.qk_nope_head_dim = mla_config.qk_nope_head_dim
    model_config.qk_rope_head_dim = mla_config.qk_rope_head_dim
    model_config.v_head_dim = mla_config.v_head_dim

    model_config.scaling = 1 / math.sqrt(
        model_config.qk_nope_head_dim + model_config.qk_rope_head_dim
    )
    rope_scaling = getattr(mla_config, "rope_scaling", None)
    if rope_scaling and "factor" in rope_scaling:
        mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
        scaling_factor = rope_scaling["factor"]
        mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
        model_config.scaling = model_config.scaling * mscale * mscale


_ATTENTION_FAMILY_SPECS = (
    _AttentionFamilySpec(
        name="DeepSeek V4",
        architectures=_DEEPSEEK_V4_ARCHITECTURES,
        configure=configure_deepseek_v4_attention,
        supports_target_verify_forward_mode=True,
        default_block_size=256,
    ),
    _AttentionFamilySpec(
        name="GLM",
        architectures=_DSA_ARCHITECTURES,
        configure=configure_glm_attention,
        default_backend="dsa",
        supports_target_verify_forward_mode=True,
    ),
    _AttentionFamilySpec(
        name="MLA",
        architectures=_MLA_ARCHITECTURES,
        configure=configure_mla_attention,
    ),
)


def _model_architectures(
    hf_config: PretrainedConfig,
    hf_text_config: PretrainedConfig,
) -> list[str]:
    return (
        [resolve_architecture(hf_config)]
        + list(getattr(hf_config, "architectures", None) or [])
        + list(getattr(hf_text_config, "architectures", []) or [])
    )


def _resolve_attention_family(
    hf_config: PretrainedConfig,
    hf_text_config: PretrainedConfig,
) -> _AttentionFamilySpec | None:
    architectures = _model_architectures(hf_config, hf_text_config)
    for spec in _ATTENTION_FAMILY_SPECS:
        if any(arch in spec.architectures for arch in architectures):
            return spec
    return None


def _apply_attention_family_defaults(
    server_args: ServerArgs,
    spec: _AttentionFamilySpec,
) -> None:
    if spec.default_block_size is not None:
        block_size_default = ServerArgs.__dataclass_fields__["block_size"].default
        if server_args.block_size == block_size_default:
            logger.info(
                "%s default block_size=%d; pass --block-size with a value other "
                "than %d to keep that value.",
                spec.name,
                spec.default_block_size,
                block_size_default,
            )
            server_args.block_size = spec.default_block_size
    if spec.default_backend is not None and server_args.attention_backend is None:
        server_args.attention_backend = spec.default_backend


def _derive_num_attention_layers(
    hf_config: PretrainedConfig,
    num_hidden_layers: int,
) -> int:
    architectures = getattr(hf_config, "architectures", None) or []
    num_attention_layers = num_hidden_layers
    if is_deepseek_v4_nextn(hf_config):
        num_attention_layers = int(getattr(hf_config, "num_nextn_predict_layers", 1))
    if any(arch in _DOUBLE_ATTENTION_LAYER_ARCHITECTURES for arch in architectures):
        num_attention_layers = num_hidden_layers * 2
    return num_attention_layers


class ModelConfig:
    def __init__(
        self,
        model_path: str,
        trust_remote_code: bool = True,
        revision: str | None = None,
        context_length: int | None = None,
        model_override_args: dict | None = None,
        dtype: str = "auto",
        quantization: str | None = None,
        override_config_file: str | None = None,
        is_draft_worker: bool | None = False,
        server_args: ServerArgs = None,
    ) -> None:
        self.model_path = model_path
        self.revision = revision
        self.quantization = quantization
        self.mapping = server_args.mapping

        # Parse args
        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()

        self.hf_config = get_config(
            model_path,
            trust_remote_code=trust_remote_code,
            revision=revision,
            model_override_args=self.model_override_args,
            is_draft_worker=is_draft_worker,
            **kwargs,
        )
        self.hf_generation_config = get_generation_config(
            self.model_path,
            trust_remote_code=trust_remote_code,
            revision=revision,
            **kwargs,
        )

        self.hf_text_config = get_hf_text_config(self.hf_config)

        # Check model type
        self.is_generation = is_generation_model(self.hf_config.architectures)
        self.is_multimodal = is_multimodal_model(self.hf_config.architectures)
        self.is_multimodal_gen = is_multimodal_gen_model(self.hf_config.architectures)
        self.is_image_gen = is_image_gen_model(self.hf_config.architectures)
        self.is_audio_model = is_audio_model(self.hf_config.architectures)

        language_model_only = bool(getattr(server_args, "language_model_only", False))
        # Target-only flag; never apply to draft / auxiliary checkpoints.
        apply_language_model_only = language_model_only and not is_draft_worker
        if apply_language_model_only:
            if not self.is_multimodal:
                raise ValueError(
                    "--language-model-only requires a multimodal model checkpoint."
                )
            logger.info(
                "Running in language-model-only mode: vision/audio encoders will "
                "be skipped; requests with multimodal inputs will be rejected."
            )
        # ``is_multimodal`` is the architectural fact; this is the runtime gate.
        self.is_multimodal_active = self.is_multimodal and not apply_language_model_only
        # Vision-only role (EPD encode): the inverse axis of language_model_only.
        # Build the vision tower (is_multimodal_active stays True) but SKIP LM
        # construction + LM weight load so a full ViT fits at encode TP=1.
        encoder_only = (
            getattr(server_args, "disaggregation_mode", None) == "encode"
            and not is_draft_worker
        )
        if encoder_only and not self.is_multimodal:
            raise ValueError(
                "disaggregation_mode=encode requires a multimodal checkpoint."
            )
        if encoder_only and apply_language_model_only:
            raise ValueError(
                "disaggregation_mode=encode (encoder-only) and language_model_only "
                "are mutually exclusive."
            )
        if encoder_only and self.is_audio_model:
            raise ValueError(
                "disaggregation_mode=encode does not support audio models; "
                "only image/video encoders are currently supported."
            )
        if encoder_only:
            # Single model-facing gate: Kimi reads hf_config.encoder_only directly;
            # Qwen3_5ForConditionalGeneration reads it to skip LM construction.
            self.hf_config.encoder_only = True
            logger.info(
                "Running in encoder-only mode: the language model will not "
                "be constructed or loaded (encode role)."
            )
        # Cap gpu_memory_utilization for VLMs in mm mode — the vision encoder
        # needs headroom that the global default doesn't account for.
        if (
            self.is_multimodal_active
            and getattr(server_args, "_gpu_memory_utilization_defaulted", False)
            and server_args.gpu_memory_utilization > 0.9
        ):
            logger.info(
                "Clamping gpu_memory_utilization %.2f -> 0.9 to leave headroom "
                "for the vision encoder.",
                server_args.gpu_memory_utilization,
            )
            server_args.gpu_memory_utilization = 0.9
        self.mm_attention_backend = getattr(server_args, "mm_attention_backend", None)
        self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)

        # Derive context length
        derived_context_len = get_context_length(self.hf_text_config)
        if context_length is not None:
            if context_length > derived_context_len:
                if envs.TOKENSPEED_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.get():
                    logger.warning(
                        "User-specified context_length (%s) is greater than the derived "
                        "context_length (%s). This may lead to incorrect model outputs or "
                        "CUDA errors.",
                        context_length,
                        derived_context_len,
                    )
                    self.context_len = context_length
                else:
                    raise ValueError(
                        f"User-specified 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. "
                        f"To allow overriding this maximum, set the env var TOKENSPEED_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1"
                    )
            else:
                self.context_len = context_length
        else:
            self.context_len = derived_context_len

        # Unify the config keys for hf_text_config
        self.head_dim = getattr(
            self.hf_text_config,
            "head_dim",
            self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads,
        )

        # MLA/DSA families carry per-head dimension metadata that does not
        # follow the standard hidden_size / num_attention_heads derivation above.
        attention_family = _resolve_attention_family(
            self.hf_config,
            self.hf_text_config,
        )
        if attention_family is not None:
            _apply_attention_family_defaults(server_args, attention_family)
            attention_family.configure(self)
        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
        else:
            self.attention_arch = AttentionArch.MHA

        self.use_v4_mtp_paged_metadata = (
            getattr(server_args, "speculative_algorithm", None) is not None
            and not is_draft_worker
            and attention_family is not None
            and attention_family.supports_target_verify_forward_mode
        )

        self.num_attention_heads = self.hf_text_config.num_attention_heads
        self.num_key_value_heads = getattr(
            self.hf_text_config, "num_key_value_heads", 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
        self.num_hidden_layers = getattr(self.hf_text_config, "num_hidden_layers", None)
        if self.num_hidden_layers is None:
            self.num_hidden_layers = self.hf_text_config.num_layers
        self.num_attention_layers = _derive_num_attention_layers(
            self.hf_config,
            self.num_hidden_layers,
        )
        if is_draft_worker:
            mtp_layers = getattr(self.hf_text_config, "mtp_num_hidden_layers", None)
            if mtp_layers is not None:
                self.num_attention_layers = mtp_layers
            else:
                nextn_layers = getattr(
                    self.hf_text_config, "num_nextn_predict_layers", None
                )
                if nextn_layers is not None and nextn_layers > 0:
                    self.num_attention_layers = nextn_layers
        self.vocab_size = self.hf_text_config.vocab_size

        # Verify quantization
        self._verify_quantization()

        # Cache attributes
        self.hf_eos_token_id = self.get_hf_eos_token_id()
        self.image_token_id = getattr(self.hf_config, "image_token_id", None)

        if server_args is not None and server_args.load_format == "extensible":
            override_model_config(self, server_args.ext_yaml)

    def _parse_quant_hf_config(self):
        quant_cfg = getattr(self.hf_config, "quantization_config", None)
        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:
            # modelopt NVFP4 checkpoints store quant config in hf_quant_config.json
            # Resolve the local snapshot directory (model_path may be a HF hub ID)
            if os.path.isdir(self.model_path):
                model_dir = self.model_path
            else:
                try:
                    from huggingface_hub import snapshot_download

                    model_dir = snapshot_download(
                        self.model_path,
                        revision=self.revision,
                        allow_patterns=["*.json"],
                        local_files_only=True,
                    )
                except Exception as exc:
                    logger.debug(
                        "Unable to resolve local quantization config for %s: %s",
                        self.model_path,
                        exc,
                    )
                    model_dir = None
            if model_dir is not None:
                hf_quant_path = os.path.join(model_dir, "hf_quant_config.json")
                if os.path.isfile(hf_quant_path):
                    with open(hf_quant_path, encoding="utf-8") as f:
                        hf_quant = json.load(f)
                    quant_algo = hf_quant.get("quantization", {}).get("quant_algo", "")
                    if quant_algo:
                        quant_cfg = {
                            "quant_method": "modelopt",
                            "quant_algo": quant_algo,
                        }
                        quant_cfg.update(hf_quant.get("quantization", {}))
        return quant_cfg

    def _verify_quantization(self) -> None:
        supported_quantization = [*QUANTIZATION_METHODS]

        optimized_quantization_methods = [
            "fp8",
            "nvfp4",
            "mxfp4",
            "compressed_tensors",
            "compressed-tensors",
            "w8a8_fp8",
        ]
        compatible_quantization_methods = {
            "w8a8_fp8": ["compressed-tensors", "compressed_tensors"],
        }
        if self.quantization is not None:
            self.quantization = self.quantization.lower()

        # Parse quantization method from the HF model config, if available.
        quant_cfg = self._parse_quant_hf_config()

        if quant_cfg is not None:
            quant_method = quant_cfg.get("quant_method", "").lower()
            # Detect which checkpoint is it
            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:
                if (
                    self.quantization not in compatible_quantization_methods
                    or quant_method
                    not in compatible_quantization_methods[self.quantization]
                ):
                    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 self.quantization not in optimized_quantization_methods:
                logger.warning(
                    "%s quantization is not fully "
                    "optimized yet. The speed can be slower than "
                    "non-quantized models.",
                    self.quantization,
                )

    def get_hf_eos_token_id(self) -> set[int] | None:
        eos_ids = getattr(self.hf_config, "eos_token_id", None)
        if eos_ids:
            # 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_hf_text_config(config: PretrainedConfig):
    """Get the "sub" config relevant to llm for multi modal models.
    No op for pure text models.
    """
    class_name = resolve_architecture(config)
    if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"):
        # We support non-hf version of llava models, so we do not want to
        # read the wrong values from the unused default text_config.
        # We set `dtype` of config to `torch.float16` for the weights, as
        # `torch.float16` is default used for image features in
        # `python/tokenspeed/runtime/models/llava.py`.
        config.dtype = torch.float16
        return config

    if hasattr(config, "thinker_config"):
        thinker_config = config.thinker_config
        if hasattr(thinker_config, "text_config"):
            return thinker_config.text_config
        return thinker_config
    if hasattr(config, "text_config"):
        if not hasattr(config.text_config, "num_attention_heads"):
            raise ValueError("text_config must define num_attention_heads")
        return config.text_config
    return config


_STR_DTYPE_TO_TORCH_DTYPE = {
    "half": torch.float16,
    "float16": torch.float16,
    "float": torch.float32,
    "float32": torch.float32,
    "bfloat16": torch.bfloat16,
}


def _get_and_verify_dtype(
    config: PretrainedConfig,
    dtype: str | torch.dtype,
) -> torch.dtype:
    # config.dtype can be missing or None.
    config_dtype = getattr(config, "dtype", None)
    if config_dtype is None:
        config_dtype = torch.bfloat16

    if isinstance(dtype, str):
        dtype = dtype.lower()
        if dtype == "auto":
            if config_dtype == torch.float32:
                if config.model_type == "gemma2":
                    logger.info(
                        "For Gemma 2, 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.info("Upcasting %s to %s.", config_dtype, torch_dtype)
        elif config_dtype == torch.float32:
            # Downcasting from float32 to float16 or bfloat16 is allowed.
            logger.info("Downcasting %s to %s.", config_dtype, torch_dtype)
        else:
            # Casting between float16 and bfloat16 is allowed with a warning.
            logger.warning("Casting %s to %s.", config_dtype, torch_dtype)

    return torch_dtype


def is_generation_model(model_architectures: list[str]):
    return True


def is_multimodal_model(model_architectures: list[str] | None):
    multimodal_architectures = {
        "Qwen3_5ForConditionalGeneration",
        "Qwen3_5MoeForConditionalGeneration",
        "Qwen3OmniMoeForConditionalGeneration",
        "Qwen3ASRForConditionalGeneration",
        "KimiK25ForConditionalGeneration",
    }
    return any(arch in multimodal_architectures for arch in model_architectures or [])


def is_multimodal_gen_model(model_architectures: list[str]):
    return False


def is_image_gen_model(model_architectures: list[str]):
    return False


def is_audio_model(model_architectures: list[str] | None):
    audio_architectures = {
        "Qwen3OmniMoeForConditionalGeneration",
        "Qwen3ASRForConditionalGeneration",
    }
    return any(arch in audio_architectures for arch in model_architectures or [])


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