modeling_qyrou_arch.py

from __future__ import annotations

from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

from .cache import QyrouArchHybridCache
from .configuration_qyrou_arch import QyrouArchConfig

_FRAMEWORK_KWARGS = frozenset(
    {
        "num_items_in_batch",
        "output_router_logits",
        "cu_seq_lens_q",
        "cu_seq_lens_k",
        "max_length_q",
        "max_length_k",
        "is_causal",
        "seq_idx",
    }
)

try:
    from .triton_kernels import PackedSwiGLUFunction
except (ImportError, RuntimeError):
    PackedSwiGLUFunction = None

try:
    from liger_kernel.ops.rms_norm import LigerRMSNormFunction
    from liger_kernel.ops.swiglu import LigerSiLUMulFunction
except ImportError:
    LigerRMSNormFunction = None
    LigerSiLUMulFunction = None

try:
    from cut_cross_entropy import linear_cross_entropy as cut_linear_cross_entropy
except ImportError:
    cut_linear_cross_entropy = None


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float, use_liger: bool = False) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
        self.use_liger = use_liger

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (
            x.is_cuda
            and self.use_liger
            and LigerRMSNormFunction is not None
            and not torch.compiler.is_compiling()
        ):
            return LigerRMSNormFunction.apply(x, self.weight, self.eps, 0.0, "llama", False, None)
        x_float = x.float()
        normalized = x_float * torch.rsqrt(x_float.square().mean(-1, keepdim=True) + self.eps)
        return normalized.to(x.dtype) * self.weight.to(x.dtype)


def _apply_rope(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> torch.Tensor:
    even, odd = x[..., 0::2], x[..., 1::2]
    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)


class CausalGQA(nn.Module):
    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.attention_backend = config.attention_backend
        if config.fused_qkv_projection:
            self.qkv_proj = nn.Linear(config.hidden_size, self.q_size + 2 * self.kv_size, bias=False)
        else:
            self.q_proj = nn.Linear(config.hidden_size, self.q_size, bias=False)
            self.k_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
            self.v_proj = nn.Linear(config.hidden_size, self.kv_size, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.o_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]
        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) if config.qk_norm else nn.Identity()

    def forward(
        self,
        x: torch.Tensor,
        rope: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        cache: QyrouArchHybridCache | None,
        cache_position: torch.Tensor,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch, query_length, _ = x.shape
        if hasattr(self, "qkv_proj"):
            q, k, v = self.qkv_proj(x).split((self.q_size, self.kv_size, self.kv_size), dim=-1)
        else:
            q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
        q = q.view(batch, query_length, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, query_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        q = _apply_rope(q, *rope)
        k = _apply_rope(k, *rope)
        if cache is not None:
            k, v = cache.update_attention(self.layer_idx, k, v, cache_position)
        key_length = k.shape[2]

        mask = None
        use_fast_causal = cache is None and attention_mask is None
        if not use_fast_causal:
            key_positions = torch.arange(key_length, device=x.device)
            allowed = key_positions[None, :] <= cache_position[:, None]
            mask = allowed[None, None, :, :].expand(batch, 1, query_length, key_length)
            if attention_mask is not None:
                if attention_mask.shape[-1] < key_length:
                    raise ValueError("attention_mask is shorter than the cached key sequence")
                mask = mask & attention_mask[:, None, None, :key_length].bool()
        if output_attentions:
            if mask is None:
                mask = torch.tril(
                    torch.ones((1, 1, query_length, key_length), dtype=torch.bool, device=x.device)
                ).expand(batch, 1, query_length, key_length)
            query_heads = q.shape[1]
            kv_heads = k.shape[1]
            if kv_heads != query_heads:
                repeat = query_heads // kv_heads
                k = k.repeat_interleave(repeat, dim=1)
                v = v.repeat_interleave(repeat, dim=1)
            scores = torch.matmul(q, k.transpose(-2, -1)) * (self.head_dim**-0.5)
            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
            weights = torch.softmax(scores, dim=-1)
            output = torch.matmul(weights, v)
        elif q.is_cuda and self.attention_backend in {"cudnn", "flash"}:
            backends = (
                [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]
                if self.attention_backend == "cudnn"
                else [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]
            )
            with sdpa_kernel(
                backends,
                set_priority=True,
            ):
                output = F.scaled_dot_product_attention(
                    q,
                    k,
                    v,
                    attn_mask=mask,
                    is_causal=use_fast_causal,
                    enable_gqa=True,
                )
        else:
            output = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=mask,
                is_causal=use_fast_causal,
                enable_gqa=True,
            )
        output = output.transpose(1, 2).contiguous().view(batch, query_length, -1)
        return self.o_proj(output), weights if output_attentions else None


class CausalConvMixer(nn.Module):
    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.kernel_size = config.conv_kernel_size
        self.depthwise = nn.Conv1d(
            config.hidden_size,
            config.hidden_size,
            kernel_size=self.kernel_size,
            groups=config.hidden_size,
            bias=False,
        )
        self.pointwise = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.pointwise._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]

    def forward(
        self,
        x: torch.Tensor,
        _rope: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        cache: QyrouArchHybridCache | None,
        cache_position: torch.Tensor,
        _output_attentions: bool = False,
    ) -> tuple[torch.Tensor, None]:
        state_length = self.kernel_size - 1
        query_mask = attention_mask[:, -x.shape[1] :] if attention_mask is not None else None
        if query_mask is not None:
            x = x * query_mask.unsqueeze(-1).to(x.dtype)
        if cache is None:
            conv_input = F.pad(x.transpose(1, 2), (state_length, 0))
        else:
            previous = cache.get_convolution(self.layer_idx)
            if previous is None:
                previous = torch.zeros(
                    (x.shape[0], state_length, x.shape[2]),
                    dtype=x.dtype,
                    device=x.device,
                )
            combined = torch.cat((previous, x), dim=1)
            conv_input = combined.transpose(1, 2)
            cache.update_convolution(self.layer_idx, combined[:, -state_length:])
        output = self.pointwise(self.depthwise(conv_input).transpose(1, 2))
        return output, None


class SwiGLU(nn.Module):
    def __init__(self, config: QyrouArchConfig) -> None:
        super().__init__()
        self.intermediate_size = config.intermediate_size
        self.use_liger = config.liger_swiglu
        if config.fused_gate_up_projection:
            self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
        else:
            self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
            self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.down_proj._qyrou_arch_residual_projection = True  # type: ignore[attr-defined]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if hasattr(self, "gate_up_proj"):
            packed = self.gate_up_proj(x)
            if (
                packed.is_cuda
                and self.use_liger
                and PackedSwiGLUFunction is not None
                and not torch.compiler.is_compiling()
            ):
                return self.down_proj(PackedSwiGLUFunction.apply(packed))
            gate, up = packed.chunk(2, dim=-1)
        else:
            gate, up = self.gate_proj(x), self.up_proj(x)
        if torch.compiler.is_compiling():
            activated = F.silu(gate) * up
        elif x.is_cuda and self.use_liger and LigerSiLUMulFunction is not None:
            activated = LigerSiLUMulFunction.apply(gate, up, 1.0, 1.0)
        else:
            activated = F.silu(gate) * up
        return self.down_proj(activated)


class QyrouArchBlock(nn.Module):
    def __init__(self, config: QyrouArchConfig, layer_idx: int) -> None:
        super().__init__()
        self.mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
        layer_number = layer_idx + 1
        self.mixer = (
            CausalConvMixer(config, layer_idx)
            if layer_number in config.conv_layers
            else CausalGQA(config, layer_idx)
        )
        self.ffn = SwiGLU(config)

    def forward(
        self,
        x: torch.Tensor,
        rope: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        cache: QyrouArchHybridCache | None,
        cache_position: torch.Tensor,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        mixer_output, weights = self.mixer(
            self.mixer_norm(x),
            rope,
            attention_mask,
            cache,
            cache_position,
            output_attentions,
        )
        x = x + mixer_output
        return x + self.ffn(self.ffn_norm(x)), weights


def _as_hybrid_cache(past_key_values: Any, num_layers: int) -> QyrouArchHybridCache:
    cache = QyrouArchHybridCache(num_layers)
    layers = getattr(past_key_values, "layers", None)
    if layers is not None:
        for index, layer in enumerate(layers[:num_layers]):
            if getattr(layer, "is_initialized", False):
                cache.attention[index] = (layer.keys, layer.values)
    else:
        key_cache = getattr(past_key_values, "key_cache", None)
        if key_cache is None:
            raise TypeError(
                "past_key_values must be a QyrouArchHybridCache or a framework Cache instance"
            )
        value_cache = past_key_values.value_cache
        for index in range(min(len(key_cache), num_layers)):
            if key_cache[index] is not None:
                cache.attention[index] = (key_cache[index], value_cache[index])
    cache.seen_tokens = past_key_values.get_seq_length()
    return cache


class QyrouArchPreTrainedModel(PreTrainedModel):
    config_class = QyrouArchConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["QyrouArchBlock"]

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)):
            scale = (
                self.config.residual_initializer_scale
                if getattr(module, "_qyrou_arch_residual_projection", False)
                else 1.0
            )
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range * scale)


class QyrouArchModel(QyrouArchPreTrainedModel):
    def __init__(self, config: QyrouArchConfig) -> None:
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [QyrouArchBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps, config.liger_rms_norm)
        head_dim = config.hidden_size // config.num_attention_heads
        inv_freq = 1.0 / (
            config.rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)
        )
        positions = torch.arange(config.max_position_embeddings, dtype=torch.float32)
        angles = torch.outer(positions, inv_freq)
        self.register_buffer("rope_cos", angles.cos(), persistent=False)
        self.register_buffer("rope_sin", angles.sin(), persistent=False)
        self.post_init()

    def _rope(
        self,
        position_ids: torch.Tensor,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if (
            not torch.compiler.is_compiling()
            and int(position_ids.max()) >= self.config.max_position_embeddings
        ):
            raise ValueError("Position exceeds max_position_embeddings")
        cos = self.rope_cos[position_ids].to(dtype).unsqueeze(1)
        sin = self.rope_sin[position_ids].to(dtype).unsqueeze(1)
        return cos, sin

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: QyrouArchHybridCache | None = None,
        inputs_embeds: torch.Tensor | None = None,
        use_cache: bool | None = None,
        cache_position: torch.LongTensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, QyrouArchHybridCache | None]:
        unsupported = set(kwargs) - _FRAMEWORK_KWARGS
        if unsupported:
            raise ValueError(f"Unsupported model arguments: {sorted(unsupported)}")
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of input_ids or inputs_embeds")
        hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
        if hidden.is_cuda and torch.is_autocast_enabled("cuda"):
            hidden = hidden.to(torch.get_autocast_dtype("cuda"))
        batch, query_length, _ = hidden.shape
        use_cache = self.config.use_cache if use_cache is None else use_cache
        return_dict = self.config.return_dict if return_dict is None else return_dict
        output_attentions = bool(
            getattr(self.config, "output_attentions", False) if output_attentions is None else output_attentions
        )
        output_hidden_states = bool(
            getattr(self.config, "output_hidden_states", False)
            if output_hidden_states is None
            else output_hidden_states
        )
        if use_cache and past_key_values is None:
            past_key_values = QyrouArchHybridCache(self.config.num_hidden_layers)
        if past_key_values is not None and not isinstance(past_key_values, QyrouArchHybridCache):
            past_key_values = _as_hybrid_cache(past_key_values, self.config.num_hidden_layers)
        cache = past_key_values if use_cache else None
        past_length = cache.get_seq_length() if cache is not None else 0
        if cache_position is None:
            cache_position = torch.arange(
                past_length,
                past_length + query_length,
                device=hidden.device,
            )
        if position_ids is None:
            if attention_mask is not None:
                position_ids = attention_mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -query_length:]
            else:
                position_ids = cache_position.unsqueeze(0).expand(batch, -1)
        rope = self._rope(position_ids, hidden.dtype)
        hidden_states = (hidden,) if output_hidden_states else None
        attentions = () if output_attentions else None
        for layer in self.layers:
            hidden, weights = layer(
                hidden,
                rope,
                attention_mask,
                cache,
                cache_position,
                output_attentions,
            )
            if hidden_states is not None:
                hidden_states += (hidden,)
            if attentions is not None:
                attentions += (weights,)
        hidden = self.norm(hidden)
        if cache is not None:
            cache.finish_step(cache_position)
        if not return_dict:
            return hidden, cache
        return BaseModelOutputWithPast(
            last_hidden_state=hidden,
            past_key_values=cache,
            hidden_states=hidden_states,
            attentions=attentions,
        )


class QyrouArchForCausalLM(QyrouArchPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: QyrouArchConfig) -> None:
        super().__init__(config)
        self.model = QyrouArchModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()
        self.tie_weights()

    def get_input_embeddings(self) -> nn.Module:
        return self.model.embed_tokens

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.model.embed_tokens = value

    def get_output_embeddings(self) -> nn.Module:
        return self.lm_head

    def set_output_embeddings(self, value: nn.Module) -> None:
        self.lm_head = value

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: QyrouArchHybridCache | None = None,
        inputs_embeds: torch.Tensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        cache_position: torch.LongTensor | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_logits: bool = True,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast | tuple[Any, ...]:
        unsupported = set(kwargs) - _FRAMEWORK_KWARGS
        if unsupported:
            raise ValueError(f"Unsupported model arguments: {sorted(unsupported)}")
        return_dict = self.config.return_dict if return_dict is None else return_dict
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            cache_position=cache_position,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,
        )
        hidden = outputs.last_hidden_state
        training_loss_only = labels is not None and not return_logits
        logits = None if training_loss_only else self.lm_head(hidden)
        loss = None
        if labels is not None:
            if (
                training_loss_only
                and hidden.is_cuda
                and self.config.cut_cross_entropy
                and cut_linear_cross_entropy is not None
            ):
                loss = cut_linear_cross_entropy(
                    hidden,
                    self.lm_head.weight,
                    labels,
                    shift=True,
                    filter_eps=None,
                )
            else:
                loss_logits = self.lm_head(hidden[:, :-1]).float()
                loss = F.cross_entropy(
                    loss_logits.reshape(-1, self.config.vocab_size),
                    labels[:, 1:].contiguous().view(-1),
                    ignore_index=-100,
                )
        if not return_dict:
            return loss, logits, outputs.past_key_values
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: QyrouArchHybridCache | None = None,
        attention_mask: torch.Tensor | None = None,
        cache_position: torch.LongTensor | None = None,
        use_cache: bool = True,
        **kwargs: Any,
    ) -> dict[str, Any]:
        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
        if past_length:
            input_ids = (
                input_ids[:, past_length:]
                if input_ids.shape[1] > past_length
                else input_ids[:, -1:]
            )
        if cache_position is None:
            cache_position = torch.arange(
                past_length,
                past_length + input_ids.shape[1],
                device=input_ids.device,
            )
        elif cache_position.numel() != input_ids.shape[1]:
            cache_position = cache_position[-input_ids.shape[1] :]
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "cache_position": cache_position,
            "use_cache": use_cache,
        }

    def _reorder_cache(
        self,
        past_key_values: QyrouArchHybridCache,
        beam_idx: torch.LongTensor,
    ) -> QyrouArchHybridCache:
        return past_key_values.reorder_cache(beam_idx)


QyrouArchConfig.register_for_auto_class()
QyrouArchForCausalLM.register_for_auto_class("AutoModelForCausalLM")