vllm-omni/vllm_omni/config/model.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from dataclasses import MISSING, field
from typing import Any

from pydantic import ConfigDict, TypeAdapter
from transformers import PretrainedConfig
from vllm.config import ModelConfig
from vllm.config.utils import config
from vllm.logger import init_logger
from vllm.transformers_utils.config import (
    get_hf_text_config,
    thinker_uses_mrope,
)
from vllm.transformers_utils.model_arch_config_convertor import (
    ModelArchConfigConvertorBase,
)

import vllm_omni.model_executor.models as me_models

logger = init_logger(__name__)

_QWEN3_TTS_TASK_TYPES = frozenset({"CustomVoice", "VoiceDesign", "Base"})


class OmniModelArchConfigConvertor(ModelArchConfigConvertorBase):
    """Config convertor for Omni multi-stage models.

    Pre-quantized checkpoints (e.g. modelopt FP8) store quantization
    config in a stage-specific sub-config (e.g.
    thinker_config.text_config.quantization_config) with correct relative
    prefixes.  The legacy hf_quant_config.json sits at the top level with
    "thinker."-prefixed names that don't match vllm-omni's module names.

    This convertor accepts an optional *stage_config_name* so that only
    the relevant stage's quantization config is surfaced.
    """

    def __init__(
        self,
        hf_config,
        hf_text_config,
        stage_config_name: str | None = None,
    ):
        super().__init__(hf_config, hf_text_config)
        self.stage_config_name = stage_config_name

    def get_quantization_config(self):
        # When a stage_config_name is set, look for quantization config
        # in that stage's text_config first (has correct relative prefixes).
        if self.stage_config_name is not None:
            stage_cfg = getattr(self.hf_config, self.stage_config_name, None)
            if stage_cfg is not None:
                text_cfg = getattr(stage_cfg, "text_config", None)
                if text_cfg is not None:
                    quant_cfg = self._normalize_quantization_config(text_cfg)
                    if quant_cfg is not None:
                        return quant_cfg

            # Fall back to top-level quantization_config
            top_quant = super().get_quantization_config()
            if top_quant is not None:
                block_names = top_quant.get("block_name_to_quantize")
                if block_names is not None:
                    # NOTE: This assumes stage_config_name follows the HF
                    # ``<stage>_config`` convention (e.g. thinker_config →
                    # prefix "thinker.").  removesuffix is a no-op when
                    # the suffix doesn't match, so a non-standard name
                    # would just use itself as prefix — safe but worth
                    # verifying if new stage names are introduced.
                    hf_prefix = self.stage_config_name.removesuffix("_config") + "."
                    if isinstance(block_names, str):
                        block_names = [b.strip() for b in block_names.split(",")]
                    if isinstance(block_names, list) and not any(b.startswith(hf_prefix) for b in block_names):
                        # This stage is not listed → no quantization.
                        return None
                return top_quant

            # For non-thinker stages (talker, code2wav) whose text_config
            # has no quantization_config, return None so quantization is
            # not applied to stages that were not quantized.
            return None

        return super().get_quantization_config()


@config(config=ConfigDict(arbitrary_types_allowed=True))
class OmniModelConfig(ModelConfig):
    """Configuration for Omni models, extending the base ModelConfig.

     This configuration class extends the base vLLM ModelConfig with
     omni-specific fields for multi-stage pipeline processing.

     Attributes:
         hf_config: The model's HF Transformers config (default: None)
         hf_text_config: The sub text_config of the model's hf_config (default: None)
         stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
         async_chunk: If set to True, perform async chunk
         session_mode: Request lifecycle mode, either turn-based or duplex
         model_stage: Stage type identifier, e.g., "thinker" or "talker"
             (default: "thinker")
         model_arch: Model architecture name
             (default: "Qwen2_5OmniForConditionalGeneration")
         worker_type: Model Type, e.g., "ar" or "generation"
         engine_output_type: Optional output type specification for the engine.
             Used to route outputs to appropriate processors (e.g., "image",
             "audio", "latent", "token_ids"). If None, output type is inferred.
         stage_connector_config: Stage connector configuration dictionary.
             Contains "name" (connector name), "extra" (extra connector config).
         task_type: Model-defined startup task type. Each model validates its
             supported values and applies the corresponding behavior.


    The correct way to initialize this class is via vLLM config, as most
    of the logic for handling values is in the ModelConfig's __post_init__.

       Example:
         >>> config = OmniModelConfig.from_vllm_model_config(
         ...     vllm_config,
         ...     stage_id=0,
         ...     model_stage="thinker",
         ...     model_arch="Qwen2_5OmniForConditionalGeneration"
         ... )
    """

    stage_id: int = 0
    # Match the upstream derived field while its value is replaced per stage.
    hf_text_config: PretrainedConfig = field(init=False)
    async_chunk: bool = False
    session_mode: str = "turn"
    retains_state_across_chunks: bool = False
    supports_running_prefix_cache_reset: bool = True
    use_v2_model_runner: bool = False
    supports_native_mrv2_data_plane: bool = False
    # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin.
    active_stream_window: int = 0
    duplex_max_sessions: int = 1
    model_stage: str = "thinker"
    model_arch: str | None = None
    worker_type: str | None = None
    engine_output_type: str | None = None
    # Optional dotted path of a per-stage pooling-output decoder applied
    # worker-side before IPC. Read by the AR scheduler.
    pooling_output_decoder: str | None = None
    final_output: bool = False
    hf_config_name: str | None = None
    custom_process_next_stage_input_func: str | None = None
    stage_connector_config: dict[str, Any] = field(
        default_factory=lambda: {
            "name": "SharedMemoryConnector",
            "extra": {},
        }
    )
    subtalker_sampling_params: dict[str, Any] | None = None
    silence_ban_frames: int = 0
    omni_kv_config: dict | None = None
    codec_frame_rate_hz: float | None = None
    task_type: str | None = None
    enable_sleep_mode: bool = False
    has_sampling_extra_args: bool = False
    # Key names (not values) of the stage's default sampling ``extra_args``.
    # Engine-core code runs before any request arrives, so this is the only
    # place it can learn which request-shaping conventions a stage uses (e.g.
    # ``cfg_role`` for classifier-free-guidance request pairs).
    sampling_extra_args_keys: tuple[str, ...] = ()
    requires_full_payload_input: bool = False

    @property
    def registry(self):
        return me_models.OmniModelRegistry

    @property
    def architectures(self) -> list[str]:
        # Falsy (None or "") means "no stage override": fall back to the
        # checkpoint config's own architectures. The stage-config builder
        # emits None; "" is tolerated for legacy callers.
        if self.model_arch:
            return [self.model_arch]
        return super().architectures

    @property
    def uses_mrope(self) -> bool:
        if self.hf_config_name is not None:
            # talker_config/thinker_config/etc
            stage_config = getattr(self.hf_config, self.hf_config_name, None)
            if stage_config is None:
                # Check the named sub-config's text_config directly.
                # Handles mrope resolution of stage-specific cls
                # (e.g., talker runs as a standalone cls)
                return thinker_uses_mrope(self.hf_config)
        return super().uses_mrope

    @property
    def embedding_size(self):
        if self.hf_config_name is not None:
            stage_config = getattr(self.hf_config, self.hf_config_name, None)
            override = getattr(stage_config, "embedding_size", None)
            if override is not None:
                return override
        return super().embedding_size

    def get_inputs_embeds_size(self) -> int:
        if self.hf_config_name is not None:
            stage_config = getattr(self.hf_config, self.hf_config_name, None)
            override = getattr(stage_config, "embedding_size", None)
            if override is not None:
                return override
        return super().get_inputs_embeds_size()

    def get_model_arch_config(self):
        # For multi-stage omni models, use a stage-aware convertor so that
        # only the correct stage's quantization config is surfaced.
        # Without this, a pre-quantized thinker checkpoint would also
        # apply quantization to the talker/code2wav stages.
        if self.hf_config_name is not None:
            convertor = OmniModelArchConfigConvertor(
                self.hf_config,
                self.hf_text_config,
                stage_config_name=self.hf_config_name,
            )
            return convertor.convert()
        return super().get_model_arch_config()

    def draw_hf_text_config(self) -> PretrainedConfig:
        # transformers' get_text_config method is used to get the text config from thinker_config.
        # to handle the case that each model stage has their own text config,
        # we need to draw the text config from the corresponding model stage.
        if self.hf_config_name is None:
            return get_hf_text_config(self.hf_config)
        try:
            # Try to get the stage-specific config (e.g., thinker_config, talker_config)
            stage_config = getattr(self.hf_config, self.hf_config_name)
            return stage_config.get_text_config()
        except AttributeError:
            # Fallback: if the attribute doesn't exist, use the default get_hf_text_config
            logger.warning(
                f"Config attribute '{self.hf_config_name}' not found in hf_config, "
                "falling back to default get_hf_text_config"
            )
            return get_hf_text_config(self.hf_config)

    def _validate_startup_task_type(self) -> None:
        """Validate startup-only task selectors owned by an AR model."""
        if self.model_arch != "Qwen3TTSTalkerForConditionalGenerationARVLLM" or self.task_type is None:
            return
        if self.task_type not in _QWEN3_TTS_TASK_TYPES:
            supported = ", ".join(sorted(_QWEN3_TTS_TASK_TYPES))
            raise ValueError(f"Qwen3-TTS --task-type must be one of {supported}; got {self.task_type!r}")

    def _patch_qwen3_tts(self):
        """Patches the value of `position_id_per_seconds` in Qwen3's
        TTS's talker_config into the this class's codec_frame_rate_hz.
        """
        talker_cfg = getattr(self.hf_config, "talker_config", None)
        if isinstance(talker_cfg, dict):
            pos_per_sec = talker_cfg.get("position_id_per_seconds")
        else:
            pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None)
        if pos_per_sec is not None:
            try:
                fps = float(pos_per_sec)
            except Exception:
                fps = None
            if fps is not None and fps > 0:
                self.codec_frame_rate_hz = fps

    def _maybe_override_text_config(self):
        """Override hf_text_config with omni-specific logic for multi-stage
        models (e.g., thinker_config, talker_config).
        """
        new_hf_text_config = self.draw_hf_text_config()
        if new_hf_text_config is not self.hf_text_config:
            self.hf_text_config = new_hf_text_config
            # Recalculate model_arch_config since it derives head counts,
            # hidden size, etc. from hf_text_config.  Without this the
            # FlashAttentionMetadataBuilder uses the wrong num_heads_q /
            # num_heads_kv (from the thinker) for talker stages, causing
            # FA3 scheduler_metadata shape mismatches at runtime.
            self.model_arch_config = self.get_model_arch_config()
            # Recalculate dependent attributes
            self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
            # Recalculate max_model_len since it depends on hf_text_config
            self.max_model_len = self.get_and_verify_max_len(self.original_max_model_len)
            # Reset sliding_window if needed
            if self.disable_sliding_window and self.hf_text_config is not None:
                self.hf_text_config.sliding_window = None

    @classmethod
    def from_vllm_model_config(cls, model_config: ModelConfig, **omni_kwargs):
        """Create OmniModelConfig from an existing vLLM ModelConfig
        and additional Omni specific kwargs.

        NOTE: The validation and __post_init__ for ModelConfig is expensive;
        to avoid calling it a second time, we explicitly retrieve defaults
        from dataclass attributes for values not passed to omni_kwargs,
        and use that to initialize a __new__ instance. This is significantly
        faster than creating the OmniModelConfig directly from the ModelConfig,
        and saves us from having to pass all kwargs to the OmniModelConfig.
        """
        # Add missing defaults to the omni kwargs and ensure values are valid
        cls.add_defaults_to_omni_kwargs(omni_kwargs)
        cls._validate_omni_fields(**omni_kwargs)

        # Allocate the new omni config and copy the model config & omni fields
        omni_cfg = object.__new__(cls)
        omni_cfg.__dict__.update(model_config.__dict__)
        omni_cfg.__dict__.update(omni_kwargs)

        # Apply any model specific patches or necessary overrides
        if (
            omni_cfg.codec_frame_rate_hz is None
            and omni_cfg.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM"
        ):
            omni_cfg._patch_qwen3_tts()

        omni_cfg._maybe_override_text_config()
        omni_cfg._validate_startup_task_type()

        if omni_cfg.hf_config is not None:
            omni_cfg.hf_config.architectures = omni_cfg.architectures

        return omni_cfg

    @classmethod
    def _validate_omni_fields(cls, **omni_kwargs):
        """Validate omni-specific fields; we use TypeAdapters here to quickly
        validate only omni kwargs to avoid rerunning validation on the
        ModelConfig.

        NOTE: This assumes add_defaults_to_omni_kwargs has already been called,
        so that all omni fields are present in the provided omni_kwargs.
        """
        omni_fields = set(cls.__dataclass_fields__) - set(ModelConfig.__dataclass_fields__)

        for key, value in omni_kwargs.items():
            if key not in omni_fields:
                raise ValueError(f"Unexpected omni kwarg: {key}")

            field_type = cls.__dataclass_fields__[key].type
            if field_type is not Any:
                TypeAdapter(field_type).validate_python(value)

        # We should not have any uninitialized keys
        uninitialized_fields = omni_fields - omni_kwargs.keys()
        if len(uninitialized_fields):
            logger.error(f"The following OmniModelConfig keys were not initialized: {uninitialized_fields}")

    @classmethod
    def add_defaults_to_omni_kwargs(cls, omni_kwargs):
        """Because we init the OmniModelConfig with __new__ to sidestep expensive
        validation, we need to be careful to ensure fields with default factories
        are initialized, otherwise we will get an AttributeError when we use it.

        To work around this issue, we explicitly add defaults to the omni_kwargs
        dict provided to ensure all fields are defined correctly.

        NOTE: omni_kwargs are mutated in place.
        """
        omni_fields = set(cls.__dataclass_fields__) - set(ModelConfig.__dataclass_fields__)

        for field_name in omni_fields - set(omni_kwargs.keys()):
            field_def = cls.__dataclass_fields__[field_name]
            if field_def.default_factory is not MISSING:
                omni_kwargs[field_name] = field_def.default_factory()
            elif field_def.default is not MISSING:
                omni_kwargs[field_name] = field_def.default