lightseekorg--tokenspeed/python/tokenspeed/runtime/models/qwen3_omni.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.

"""Inference-only Qwen3-Omni thinker (text output only).

The thinker composes TokenSpeed's shared Qwen3 MoE language model, Qwen3 vision
tower, and Qwen3-Omni audio tower. Talker/Code2Wav weights are deliberately
ignored: this entry point implements ``return_audio=False`` and independent
image/audio/video inputs (``use_audio_in_video=False``).
"""

from __future__ import annotations

import logging
from collections.abc import Iterable

import torch

from tokenspeed.runtime.configs.qwen3_vision_config import Qwen3VLVisionConfig
from tokenspeed.runtime.distributed.mapping import Mapping
from tokenspeed.runtime.execution.context import ForwardContext
from tokenspeed.runtime.layers.logits_processor import LogitsMetadata
from tokenspeed.runtime.layers.moe import (
    ExpertCheckpointSchema,
    build_moe_checkpoint_loader,
)
from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
from tokenspeed.runtime.layers.utils import get_layer_id
from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader
from tokenspeed.runtime.models.qwen3_audio import Qwen3AudioEncoder
from tokenspeed.runtime.models.qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel
from tokenspeed.runtime.models.qwen3_vision import Qwen3VLMoeVisionModel
from tokenspeed.runtime.multimodal.embedder import (
    EncoderSpec,
    MultimodalEmbedder,
    pad_input_tokens,
)
from tokenspeed.runtime.multimodal.encoder_cudagraph import (
    EncoderCudaGraphWrapper,
    VisionEncoderCudaGraphAdapter,
)
from tokenspeed.runtime.multimodal.inputs import (
    Modality,
    MultimodalDataItem,
    MultimodalInputs,
)
from tokenspeed.runtime.utils.env import envs

logger = logging.getLogger(__name__)


def _get_thinker_config(config):
    return getattr(config, "thinker_config", config)


def _shared_vision_config(vision_config) -> Qwen3VLVisionConfig:
    if not getattr(vision_config, "apply_vit_abs_pos_embed", True):
        raise ValueError("Qwen3-Omni without absolute vision positions is unsupported")
    values = (
        vision_config.to_dict()
        if hasattr(vision_config, "to_dict")
        else dict(vars(vision_config))
    )
    patch_size = int(values.get("patch_size", 16))
    image_size = values.get("image_size")
    if image_size is not None:
        image_size = int(image_size)
        if image_size % patch_size:
            raise ValueError("Qwen3-Omni image_size must be divisible by patch_size")
        expected_positions = (image_size // patch_size) ** 2
        configured_positions = values.get("num_position_embeddings")
        if (
            configured_positions is not None
            and int(configured_positions) != expected_positions
        ):
            raise ValueError(
                "Qwen3-Omni image_size/patch_size disagrees with "
                "num_position_embeddings"
            )
        values["num_position_embeddings"] = expected_positions
    values["deepstack_visual_indexes"] = (
        values.get("deepstack_visual_indexes", [8, 16, 24]) or []
    )
    return Qwen3VLVisionConfig(**values)


class Qwen3OmniMoeTextModel(Qwen3MoeModel):
    """Qwen3 MoE decoder with Omni visual deepstack injection."""

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        ctx: ForwardContext,
        out_cache_loc: torch.Tensor,
        input_embeds: torch.Tensor | None = None,
        input_deepstack_embeds: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, None]:
        hidden_states = (
            self.embed_tokens(input_ids) if input_embeds is None else input_embeds
        )
        residual = None
        hidden_size = self.config.hidden_size
        num_deepstack = (
            input_deepstack_embeds.shape[-1] // hidden_size
            if input_deepstack_embeds is not None
            else 0
        )

        for layer_idx, layer in enumerate(self.layers):
            hidden_states, residual = layer(
                positions,
                hidden_states,
                ctx,
                out_cache_loc,
                residual,
                cos_sin=None,
            )
            if layer_idx < num_deepstack and input_deepstack_embeds.numel() > 0:
                start = layer_idx * hidden_size
                hidden_states.add_(
                    input_deepstack_embeds[:, start : start + hidden_size]
                )

        if not ctx.forward_mode.is_idle():
            hidden_states, _ = layer.comm_manager.final_norm(
                hidden_states, residual, ctx, self.norm
            )
        return hidden_states, None


class Qwen3OmniMoeForConditionalGeneration(Qwen3MoeForCausalLM):
    """Qwen3-Omni thinker for text-only generation."""

    model_cls = Qwen3OmniMoeTextModel

    def __init__(
        self,
        config,
        mapping: Mapping,
        quant_config: QuantizationConfig | None = None,
        is_multimodal_active: bool = True,
        mm_attention_backend: str | None = None,
    ) -> None:
        self.omni_config = config
        self.thinker_config = _get_thinker_config(config)
        text_config = self.thinker_config.text_config
        super().__init__(config=text_config, mapping=mapping, quant_config=quant_config)

        self.is_multimodal_active = is_multimodal_active
        self.multimodal_embedder = (
            MultimodalEmbedder() if is_multimodal_active else None
        )
        if not is_multimodal_active:
            self.visual = None
            self.audio_tower = None
            self.deepstack_visual_indexes = []
            self.num_deepstack_embeddings = 0
            self.image_encoder = None
            self.video_encoder = None
            self.audio_encoder = None
            return

        vision_config = _shared_vision_config(self.thinker_config.vision_config)
        if vision_config.out_hidden_size != text_config.hidden_size:
            raise ValueError(
                "Qwen3-Omni vision output size must match text hidden size"
            )
        if self.thinker_config.audio_config.output_dim != text_config.hidden_size:
            raise ValueError("Qwen3-Omni audio output size must match text hidden size")

        self.visual = Qwen3VLMoeVisionModel(
            vision_config,
            mapping=mapping,
            quant_config=None,
            norm_eps=getattr(text_config, "rms_norm_eps", 1e-6),
            prefix="visual",
            mm_attention_backend=mm_attention_backend,
        )
        self.audio_tower = Qwen3AudioEncoder(
            self.thinker_config.audio_config,
            mapping=mapping,
            quant_config=None,
            prefix="audio_tower",
            mm_attention_backend=mm_attention_backend,
        )
        self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes
        self.num_deepstack_embeddings = len(self.deepstack_visual_indexes)

        # Encoder callables can be replaced by CUDA graph wrappers at runtime.
        self.image_encoder = self.get_image_feature
        self.video_encoder = self.get_video_feature
        self.audio_encoder = self.audio_tower.encode

    def separate_deepstack_embeds(
        self, embedding: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        expected_parts = 1 + self.num_deepstack_embeddings
        if embedding.shape[-1] != self.config.hidden_size * expected_parts:
            raise ValueError(
                f"vision embedding width {embedding.shape[-1]} does not match "
                f"{expected_parts} x text hidden size {self.config.hidden_size}"
            )
        split = self.config.hidden_size
        return embedding[:, :split], embedding[:, split:]

    def pad_input_ids(
        self, input_ids: list[int], mm_inputs: MultimodalInputs
    ) -> list[int]:
        return pad_input_tokens(input_ids, mm_inputs)

    def pre_encode(
        self, items: list[MultimodalDataItem]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        pixel_values = torch.cat([item.feature for item in items], dim=0).type(
            self.visual.dtype
        )
        grid = torch.cat(
            [
                (
                    item.video_grid_thw
                    if item.modality == Modality.VIDEO
                    else item.image_grid_thw
                )
                for item in items
            ],
            dim=0,
        )
        if pixel_values.dim() != 2 or grid.dim() != 2:
            raise ValueError("Qwen3-Omni vision features require 2-D patches and grids")
        return self.visual.prepare_patch_embed(pixel_values, grid), grid

    def post_encode(
        self, encoder_outs: list[torch.Tensor], grid: torch.Tensor
    ) -> torch.Tensor:
        del grid
        return torch.cat(encoder_outs, dim=0)

    def get_image_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor:
        tokens, grid = self.pre_encode(items)
        output = self.visual.forward_blocks(tokens, self.visual.prepare_metadata(grid))
        return self.post_encode([output], grid)

    def get_video_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor:
        tokens, grid = self.pre_encode(items)
        output = self.visual.forward_blocks(tokens, self.visual.prepare_metadata(grid))
        return self.post_encode([output], grid)

    def _build_encoder_cudagraph_wrapper(
        self,
        mapping: Mapping,
        *,
        max_metadata_sequences_per_batch: int | None = None,
        metadata_sequence_budget_from_encoder_output_budget: bool = False,
    ) -> EncoderCudaGraphWrapper:
        adapter = VisionEncoderCudaGraphAdapter(
            tower=self.visual,
            pre_encode=self.pre_encode,
            post_encode=self.post_encode,
            out_div=self.visual.spatial_merge_size**2,
            merge=self.visual.spatial_merge_size,
            input_feature_shape=(1, self.visual.hidden_size),
            modality_name="vision",
            capture_tp_size=mapping.vision.tp_size,
            capture_tp_group=mapping.vision.tp_group,
        )
        return EncoderCudaGraphWrapper(
            adapter=adapter,
            budget_range=(64, 4096),
            max_metadata_sequences_per_batch=max_metadata_sequences_per_batch,
            metadata_sequence_budget_from_encoder_output_budget=(
                metadata_sequence_budget_from_encoder_output_budget
            ),
        )

    def make_encoder_cudagraph_wrappers(self, mapping: Mapping) -> dict:
        max_video_sequences = (
            envs.TOKENSPEED_MM_VIDEO_ENCODER_CUDA_GRAPH_MAX_SEQUENCES_PER_BATCH.get()
        )
        if max_video_sequences is not None:
            max_video_sequences = max(1, max_video_sequences)
        shared = self._build_encoder_cudagraph_wrapper(
            mapping,
            max_metadata_sequences_per_batch=max_video_sequences,
            metadata_sequence_budget_from_encoder_output_budget=(
                max_video_sequences is None
            ),
        )
        return {"image_encoder": shared, "video_encoder": shared}

    @torch.no_grad()
    def forward(
        self,
        ctx: ForwardContext,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        out_cache_loc: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        multimodal_context = kwargs.pop("multimodal_context", None)
        if (
            not self.is_multimodal_active
            and multimodal_context is not None
            and multimodal_context.has_extend_inputs()
        ):
            raise RuntimeError(
                "Qwen3-Omni received multimodal inputs while its encoders are disabled"
            )
        if (
            multimodal_context is None
            or not multimodal_context.has_extend_inputs()
            or ctx.forward_mode.is_decode_or_idle()
        ):
            return super().forward(ctx, input_ids, positions, out_cache_loc, **kwargs)

        input_embeds, model_kwargs = self.multimodal_embedder.apply(
            input_ids=input_ids,
            text_embedding=self.model.embed_tokens,
            ctx=multimodal_context,
            encoders={
                Modality.IMAGE: EncoderSpec(self.image_encoder, deepstack=True),
                Modality.VIDEO: EncoderSpec(self.video_encoder, deepstack=True),
                Modality.AUDIO: EncoderSpec(self.audio_encoder),
            },
            multimodal_model=self,
            is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(),
        )
        hidden_states, aux_hidden_states = self.model(
            input_ids,
            positions,
            ctx,
            out_cache_loc,
            input_embeds=input_embeds,
            **model_kwargs,
        )
        return self.logits_processor(
            input_ids,
            hidden_states,
            self.lm_head,
            LogitsMetadata.from_forward_context(ctx),
            aux_hidden_states,
        )

    @staticmethod
    def _map_visual_weight(name: str) -> str:
        name = name.replace("attn.qkv.", "attn.qkv_proj.")
        name = name.replace("merger_list.", "deepstack_merger_list.")
        name = name.replace(".ln_q.", ".norm.")
        name = name.replace(".mlp.0.", ".linear_fc1.")
        name = name.replace(".mlp.2.", ".linear_fc2.")
        return name

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        stacked_params_mapping = [
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        ignore_suffixes = (
            ".bias",
            "_bias",
            ".k_scale",
            "_k_scale",
            ".v_scale",
            "_v_scale",
            ".weight_scale",
            "_weight_scale",
            ".input_scale",
            "_input_scale",
        )
        params = dict(self.named_parameters(remove_duplicate=False))
        moe_loader = build_moe_checkpoint_loader(
            params_dict=params,
            expert_schema=ExpertCheckpointSchema(
                gate_proj_name="gate_proj",
                down_proj_name="down_proj",
                up_proj_name="up_proj",
            ),
            fused_schema=ExpertCheckpointSchema(
                gate_up_fused_name="gate_up_proj",
                down_proj_name="down_proj",
            ),
            num_experts=self.config.num_experts,
            ep_rank=self.mapping.moe.ep_rank,
            ep_size=self.mapping.moe.ep_size,
        )
        loaded = set()

        for original_name, loaded_weight in weights:
            if original_name.startswith(("talker.", "code2wav.")):
                continue
            name = (
                original_name[len("thinker.") :]
                if original_name.startswith("thinker.")
                else original_name
            )

            if name.startswith("visual."):
                if not self.is_multimodal_active:
                    continue
                name = self._map_visual_weight(name)
                if name not in params:
                    logger.warning("Parameter %s not found in Qwen3-Omni", name)
                    continue
                param = params[name]
                loader = getattr(param, "weight_loader", default_weight_loader)
                loader(param, loaded_weight)
                loaded.add(name)
                continue

            if name.startswith("audio_tower."):
                if not self.is_multimodal_active:
                    continue
                loaded_name = self.audio_tower.load_weight(name, loaded_weight)
                if loaded_name is None:
                    logger.warning("Parameter %s not found in Qwen3-Omni", name)
                else:
                    loaded.add(f"audio_tower.{loaded_name}")
                continue

            layer_id = get_layer_id(name)
            if (
                layer_id is not None
                and hasattr(self.model, "start_layer")
                and not self.model.start_layer <= layer_id < self.model.end_layer
            ):
                continue
            if "rotary_emb" in name:
                continue
            if self.config.tie_word_embeddings and name == "lm_head.weight":
                continue

            for param_name, weight_name, shard_id in stacked_params_mapping:
                if weight_name not in name or "mlp.experts" in name:
                    continue
                mapped_name = name.replace(weight_name, param_name)
                if mapped_name.endswith(ignore_suffixes) and mapped_name not in params:
                    break
                if mapped_name not in params:
                    break
                param = params[mapped_name]
                param.weight_loader(param, loaded_weight, shard_id)
                loaded.add(mapped_name)
                break
            else:
                if name.endswith((".bias", "_bias")) and name not in params:
                    continue
                if moe_loader.matches(name):
                    loaded.add(moe_loader.load(name, loaded_weight))
                    continue
                if moe_loader.is_expert_checkpoint_weight(name):
                    continue
                if name.endswith(ignore_suffixes) and name not in params:
                    continue
                if name not in params:
                    logger.warning("Parameter %s not found in Qwen3-Omni", name)
                    continue
                param = params[name]
                loader = getattr(param, "weight_loader", default_weight_loader)
                loader(param, loaded_weight)
                loaded.add(name)

        return loaded


EntryClass = [Qwen3OmniMoeForConditionalGeneration]