vllm-omni/vllm_omni/model_executor/models/personaplex/modeling_helium.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""vLLM-native PersonaPlex Helium temporal transformer.

The module maps the Moshi temporal LM backbone onto vLLM's Llama-style decoder
components. It intentionally does not include Moshi's text/audio input
embeddings or depformer; PersonaPlex feeds temporal `inputs_embeds` directly.
"""

from __future__ import annotations

from collections.abc import Iterable, Mapping
from itertools import islice

import torch
from torch import nn
from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, VllmConfig
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
from vllm.model_executor.models.utils import (
    PPMissingLayer,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
from vllm.sequence import IntermediateTensors
from vllm.v1.attention.backend import AttentionType

from vllm_omni.model_executor.models.personaplex.configuration_helium import (
    HeliumConfig,
)

__all__ = ["HeliumForCausalLM", "HeliumModel"]


class HeliumRMSNorm(RMSNorm):
    """Moshi-compatible RMSNorm.

    Moshi stores the parameter as `alpha` with shape `[1, 1, hidden]` and
    computes variance and the alpha multiply in fp32 before casting back to the
    input dtype. vLLM's parameter is named `weight`; the loader squeezes alpha.
    """

    def _forward_f32(self, x: torch.Tensor) -> torch.Tensor:
        output_dtype = x.dtype
        x_f32 = x.float()
        variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
        x_f32 = x_f32 * torch.rsqrt(variance + self.variance_epsilon)
        if self.has_weight:
            x_f32 = x_f32 * self.weight.float()
        return x_f32.to(output_dtype)

    def forward(
        self,
        x: torch.Tensor,
        residual: torch.Tensor | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            return self._forward_f32(x)
        residual = residual + x
        return self._forward_f32(residual), residual


class HeliumMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.")
        self.gate_up_proj = MergedColumnParallelLinear(
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
        )
        self.act_fn = SiluAndMul()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
        x, _ = self.down_proj(x)
        return x


def _get_helium_rope(config: HeliumConfig):
    rope_kwargs = {
        "head_size": config.head_dim,
        "max_position": config.max_position_embeddings,
        "is_neox_style": False,
        "rope_parameters": {"base": config.rope_theta},
    }
    try:
        return get_rope(**rope_kwargs)
    except TypeError:
        # vLLM 0.23 uses `rope_parameters`; this fallback is for older local
        # source trees whose get_rope still takes rotary_dim/base/rope_scaling.
        return get_rope(
            head_size=config.head_dim,
            rotary_dim=config.head_dim,
            max_position=config.max_position_embeddings,
            base=config.rope_theta,
            is_neox_style=False,
        )


class HeliumAttention(nn.Module):
    def __init__(
        self,
        config: HeliumConfig,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        if self.total_num_kv_heads >= tp_size:
            assert self.total_num_kv_heads % tp_size == 0
        else:
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = config.head_dim
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5

        self.qkv_proj = QKVParallelLinear(
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.qkv_proj",
        )
        self.o_proj = RowParallelLinear(
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )
        self.rotary_emb = _get_helium_rope(config)
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            per_layer_sliding_window=config.sliding_window,
            prefix=f"{prefix}.attn",
            attn_type=attn_type,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output


class HeliumDecoderLayer(nn.Module):
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        config: HeliumConfig | None = None,
    ) -> None:
        super().__init__()
        config = config or vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        self.hidden_size = config.hidden_size

        self.self_attn = HeliumAttention(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
            attn_type=AttentionType.DECODER,
        )
        self.mlp = HeliumMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
        self.input_layernorm = HeliumRMSNorm(
            config.hidden_size,
            eps=config.rms_norm_eps,
        )
        self.post_attention_layernorm = HeliumRMSNorm(
            config.hidden_size,
            eps=config.rms_norm_eps,
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)

        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        hidden_states, residual = self.post_attention_layernorm(
            hidden_states,
            residual,
        )
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


def _flatten_inputs(
    hidden_states: torch.Tensor,
    positions: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int] | None]:
    if hidden_states.dim() != 3:
        return hidden_states, positions.reshape(-1), None

    batch, seq_len, hidden = hidden_states.shape
    flat_hidden = hidden_states.reshape(batch * seq_len, hidden)
    if positions.dim() == 1:
        if positions.numel() == seq_len:
            flat_positions = positions.unsqueeze(0).expand(batch, seq_len).reshape(-1)
        elif positions.numel() == batch * seq_len:
            flat_positions = positions.reshape(-1)
        else:
            raise ValueError(
                "1D positions must have length S or B*S when inputs_embeds is "
                f"3D, got {positions.numel()} for B={batch}, S={seq_len}."
            )
    else:
        flat_positions = positions.reshape(-1)
    return flat_hidden, flat_positions, (batch, seq_len)


@support_torch_compile(
    dynamic_arg_dims={
        "input_ids": 0,
        "positions": 0,
        "intermediate_tensors": 0,
        "inputs_embeds": 0,
    }
)
class HeliumModel(nn.Module):
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = HeliumDecoderLayer,
        config: HeliumConfig | None = None,
    ) -> None:
        super().__init__()
        # ``config`` lets a composite model (e.g. the PersonaPlex talker) build the
        # Helium backbone from its temporal sub-config while the engine's top-level
        # hf_config is the composite PersonaPlexConfig. Defaults to the engine's.
        config = config or vllm_config.model_config.hf_config
        self.config = config
        self.quant_config = vllm_config.quant_config
        self.vocab_size = config.vocab_size

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix, config=config),
            prefix=f"{prefix}.layers",
        )
        if get_pp_group().is_last_rank:
            self.norm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()

        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"],
            config.hidden_size,
        )

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        raise ValueError(
            "Helium temporal forward requires inputs_embeds; Moshi token "
            "embeddings are intentionally not part of this vLLM module."
        )

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
        unflatten_shape: tuple[int, int] | None = None
        if get_pp_group().is_first_rank:
            if inputs_embeds is None:
                raise ValueError("HeliumModel.forward requires inputs_embeds.")
            hidden_states, positions, unflatten_shape = _flatten_inputs(
                inputs_embeds,
                positions,
            )
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
            positions = positions.reshape(-1)

        for layer in islice(self.layers, self.start_layer, self.end_layer):
            hidden_states, residual = layer(
                positions,
                hidden_states,
                residual,
            )

        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})

        hidden_states, _ = self.norm(hidden_states, residual)
        if unflatten_shape is not None:
            batch, seq_len = unflatten_shape
            hidden_states = hidden_states.reshape(batch, seq_len, -1)
        return hidden_states


class HeliumForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
    embedding_modules = {
        "lm_head": "output_embeddings",
    }

    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = HeliumDecoderLayer,
    ) -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        self.config = config
        self.quant_config = vllm_config.quant_config
        self.model = HeliumModel(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "model"),
            layer_type=layer_type,
        )
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=self.quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
        else:
            self.lm_head = PPMissingLayer()
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)

    def forward(
        self,
        input_ids: torch.Tensor | None,
        positions: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
        return self.model(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
        if hidden_states.dim() == 3:
            batch, seq_len, hidden = hidden_states.shape
            flat_hidden = hidden_states.reshape(batch * seq_len, hidden)
            logits = self.logits_processor(self.lm_head, flat_hidden)
            if logits is None:
                return None
            return logits.reshape(batch, seq_len, -1)
        return self.logits_processor(self.lm_head, hidden_states)

    def load_weights(
        self,
        weights: Iterable[tuple[str, torch.Tensor]] | Mapping[str, torch.Tensor],
    ) -> set[str]:
        if isinstance(weights, Mapping):
            weights = weights.items()

        params_dict = dict(self.named_parameters(remove_duplicate=False))
        loaded_params: set[str] = set()

        for name, loaded_weight in weights:
            mapped = self._load_moshi_weight(name, loaded_weight, params_dict)
            loaded_params.update(mapped)

        return loaded_params

    def _load_moshi_weight(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        params_dict: dict[str, nn.Parameter],
    ) -> set[str]:
        if name == "out_norm.alpha":
            return self._load_direct(
                "model.norm.weight",
                loaded_weight.squeeze(),
                params_dict,
            )
        if name == "text_linear.weight":
            return self._load_direct("lm_head.weight", loaded_weight, params_dict)

        prefix = "transformer.layers."
        if not name.startswith(prefix):
            return set()

        rest = name.removeprefix(prefix)
        layer_index, _, suffix = rest.partition(".")
        if not layer_index.isdigit() or not suffix:
            return set()

        base = f"model.layers.{layer_index}"
        if suffix == "self_attn.in_proj_weight":
            q_weight, k_weight, v_weight = loaded_weight.chunk(3, dim=0)
            param_name = f"{base}.self_attn.qkv_proj.weight"
            self._load_shard(param_name, q_weight, "q", params_dict)
            self._load_shard(param_name, k_weight, "k", params_dict)
            self._load_shard(param_name, v_weight, "v", params_dict)
            return {param_name}
        if suffix == "self_attn.out_proj.weight":
            return self._load_direct(
                f"{base}.self_attn.o_proj.weight",
                loaded_weight,
                params_dict,
            )
        if suffix == "gating.linear_in.weight":
            gate_weight, up_weight = loaded_weight.chunk(2, dim=0)
            param_name = f"{base}.mlp.gate_up_proj.weight"
            self._load_shard(param_name, gate_weight, 0, params_dict)
            self._load_shard(param_name, up_weight, 1, params_dict)
            return {param_name}
        if suffix == "gating.linear_out.weight":
            return self._load_direct(
                f"{base}.mlp.down_proj.weight",
                loaded_weight,
                params_dict,
            )
        if suffix == "norm1.alpha":
            return self._load_direct(
                f"{base}.input_layernorm.weight",
                loaded_weight.squeeze(),
                params_dict,
            )
        if suffix == "norm2.alpha":
            return self._load_direct(
                f"{base}.post_attention_layernorm.weight",
                loaded_weight.squeeze(),
                params_dict,
            )
        return set()

    def _load_direct(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        params_dict: dict[str, nn.Parameter],
    ) -> set[str]:
        if name not in params_dict or is_pp_missing_parameter(name, self):
            return set()
        param = params_dict[name]
        weight_loader = getattr(param, "weight_loader", default_weight_loader)
        weight_loader(param, loaded_weight)
        return {name}

    def _load_shard(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        shard_id: str | int,
        params_dict: dict[str, nn.Parameter],
    ) -> None:
        if name not in params_dict or is_pp_missing_parameter(name, self):
            return
        param = params_dict[name]
        weight_loader = param.weight_loader
        weight_loader(param, loaded_weight, shard_id)