# AetherMind — modeling_aethermind.py
# Copyright 2026 AetherMind Project. Apache License 2.0.
"""PyTorch AetherMind model: a modern decoder-only Transformer.

Components implemented from scratch (no third-party LLM code):

    * ``AetherMindRMSNorm``          — RMSNorm with fp32 statistics
    * ``AetherMindRotaryEmbedding``  — RoPE with optional linear scaling
    * ``AetherMindAttention``        — Grouped-Query Attention + KV cache
    * ``AetherMindMLP``              — SwiGLU feed-forward
    * ``AetherMindSparseMoE``        — top-k Mixture-of-Experts + aux loss
    * ``AetherMindDecoderLayer``     — pre-norm Transformer block
    * ``AetherMindModel``            — embedding + blocks + final norm
    * ``AetherMindForCausalLM``      — LM head + loss + generation

Attention runs through ``torch.nn.functional.scaled_dot_product_attention``:
on CUDA this dispatches to FlashAttention kernels when no explicit mask is
needed (no padding), and to the memory-efficient backend otherwise.
"""

from __future__ import annotations

import math
import warnings
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GenerationMixin
from transformers.cache_utils import DynamicCache
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .configuration_aethermind import AetherMindConfig

logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "AetherMindConfig"


# ======================================================================
# Building blocks
# ======================================================================
class AetherMindRMSNorm(nn.Module):
    """Root-mean-square layer normalization (as in the Llama/RWKV family)."""

    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


class AetherMindRotaryEmbedding(nn.Module):
    """Rotary position embeddings with optional linear interpolation scaling."""

    def __init__(
        self,
        head_dim: int,
        max_position_embeddings: int = 4096,
        base: float = 10000.0,
        rope_scaling: Optional[dict] = None,
    ) -> None:
        super().__init__()
        self.head_dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        factor = 1.0
        if rope_scaling is not None:
            rtype = rope_scaling.get("rope_type", rope_scaling.get("type", "linear"))
            if rtype == "linear":
                factor = float(rope_scaling.get("factor", 1.0))
        self.scaling_factor = factor
        # inv_freq is kept as a plain attribute (NOT a registered buffer):
        # it is deterministic from (head_dim, base, scaling_factor), so saving it
        # is redundant; keeping it out of the state dict avoids meta-device
        # materialization bugs and accidental dtype casts (e.g. model.to(bf16)).
        self._inv_freq: Optional[torch.Tensor] = None

    def _get_inv_freq(self, device: torch.device) -> torch.Tensor:
        if (
            self._inv_freq is None
            or self._inv_freq.device != device
            or self._inv_freq.dtype != torch.float32
        ):
            inv_freq = 1.0 / (
                self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim)
            )
            self._inv_freq = inv_freq / self.scaling_factor
        return self._inv_freq

    @torch.no_grad()
    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
        """Return cos/sin of shape ``(batch, seq_len, head_dim)`` in float32."""
        inv_freq = self._get_inv_freq(x.device)[None, None, :]        # (1, 1, hd/2)
        pos = position_ids[:, :, None].to(torch.float32)              # (b, s, 1)
        freqs = pos * inv_freq                                        # (b, s, hd/2)
        emb = torch.cat((freqs, freqs), dim=-1)                       # (b, s, hd)
        if position_ids.max() >= self.max_position_embeddings:
            warnings.warn(
                "Sequence length exceeds max_position_embeddings; RoPE positions "
                "beyond the trained window degrade quality. Consider context extension."
            )
        return emb.cos().to(torch.float32), emb.sin().to(torch.float32)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim: int = 1):
    cos = cos.unsqueeze(unsqueeze_dim)  # (b, 1, s, d)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Expand KV heads to match the number of query heads (GQA)."""
    if n_rep == 1:
        return hidden_states
    batch, num_kv_heads, slen, head_dim = hidden_states.shape
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_kv_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)


# ======================================================================
# Attention
# ======================================================================
class AetherMindAttention(nn.Module):
    """Grouped-Query Attention with RoPE and a dynamic KV cache."""

    def __init__(self, config: AetherMindConfig, layer_idx: int) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = config.head_dim or config.hidden_size // config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        if self.num_heads % self.num_key_value_heads != 0:
            raise ValueError(
                f"num_attention_heads ({self.num_heads}) must be divisible by "
                f"num_key_value_heads ({self.num_key_value_heads})"
            )
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.scaling = self.head_dim ** -0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True

        op_size = self.num_heads * self.head_dim
        self.q_proj = nn.Linear(self.hidden_size, op_size, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(op_size, self.hidden_size, bias=config.attention_bias)

        self.rotary_emb = AetherMindRotaryEmbedding(
            self.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
            config.rope_scaling,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        use_cache: bool = False,
        cache_position: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        bsz, q_len, _ = hidden_states.shape

        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_values is not None:
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_values.update(
                key_states, value_states, self.layer_idx, cache_kwargs
            )

        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        if output_attentions:
            # Explicit path (returns attention probabilities, CPU-friendly)
            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling
            if attention_mask is not None:
                attn_weights = attn_weights + attention_mask
            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
            attn_output = torch.matmul(attn_weights, value_states)
            attn_weights_out = attn_weights
        else:
            full_prefill = attention_mask is None and q_len == key_states.shape[-2] and q_len > 1
            attn_output = F.scaled_dot_product_attention(
                query_states,
                key_states,
                value_states,
                attn_mask=None if full_prefill else attention_mask,
                dropout_p=self.attention_dropout if self.training else 0.0,
                is_causal=full_prefill,
                scale=self.scaling,
            )
            attn_weights_out = None

        attn_output = attn_output.transpose(1, 2).contiguous().reshape(bsz, q_len, -1)
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights_out


# ======================================================================
# Feed-forward blocks
# ======================================================================
ACT2FN = {"silu": F.silu, "gelu": F.gelu, "relu": F.relu}


class AetherMindMLP(nn.Module):
    """SwiGLU feed-forward: down( act(gate(x)) * up(x) )."""

    def __init__(self, config: AetherMindConfig, intermediate_size: Optional[int] = None) -> None:
        super().__init__()
        inter = intermediate_size or config.intermediate_size
        self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, inter, bias=False)
        self.down_proj = nn.Linear(inter, config.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class AetherMindSparseMoE(nn.Module):
    """Top-k Mixture-of-Experts feed-forward with a load-balancing aux loss.

    The router picks ``num_experts_per_tok`` experts per token; outputs are
    the probability-weighted sum of the selected experts. A switch-transformer
    style auxiliary loss (importance x load) encourages balanced routing and
    is accumulated in ``AetherMindModel.moe_aux_loss`` during training.
    """

    def __init__(self, config: AetherMindConfig) -> None:
        super().__init__()
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_tok
        self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
        self.experts = nn.ModuleList(
            [AetherMindMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
        )

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        bsz, seq_len, hidden = x.shape
        flat = x.reshape(-1, hidden)
        router_logits = self.gate(flat)                                   # (N, E)
        routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32)
        topk_weights, topk_ids = torch.topk(routing_probs, self.top_k, dim=-1)
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)

        out = torch.zeros_like(flat)
        for e_idx, expert in enumerate(self.experts):
            mask = topk_ids == e_idx                                      # (N, k)
            token_mask = mask.any(dim=-1)
            if token_mask.any():
                weights = (topk_weights * mask).sum(dim=-1)[token_mask].to(flat.dtype)
                out[token_mask] += expert(flat[token_mask]) * weights.unsqueeze(-1)

        # Auxiliary load-balancing loss (importance x load), switch-transformer style
        importance = routing_probs.mean(dim=0)                            # (E,)
        load = topk_ids.new_zeros(self.num_experts, dtype=torch.float32)
        for k in range(self.top_k):
            load += torch.bincount(topk_ids[:, k], minlength=self.num_experts).to(torch.float32)
        load = load / (flat.shape[0] * self.top_k)
        aux_loss = self.num_experts * torch.sum(importance * load)

        return out.reshape(bsz, seq_len, hidden), aux_loss


# ======================================================================
# Decoder layer
# ======================================================================
class AetherMindDecoderLayer(nn.Module):
    """Pre-norm block: x + Attn(RMSNorm(x)); x + FFN(RMSNorm(x))."""

    def __init__(self, config: AetherMindConfig, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.self_attn = AetherMindAttention(config, layer_idx=layer_idx)
        use_moe = config.use_moe and (layer_idx % max(1, config.moe_layers_freq) == 0)
        self.mlp = AetherMindSparseMoE(config) if use_moe else AetherMindMLP(config)
        self.input_layernorm = AetherMindRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = AetherMindRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[DynamicCache] = None,
        use_cache: bool = False,
        cache_position: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        attn_out, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            output_attentions=output_attentions,
        )
        hidden_states = residual + attn_out

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        aux_loss = None
        if isinstance(self.mlp, AetherMindSparseMoE):
            hidden_states, aux_loss = self.mlp(hidden_states)
        else:
            hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states, attn_weights, aux_loss)
        return outputs


# ======================================================================
# Base model
# ======================================================================
class AetherMindPreTrainedModel(PreTrainedModel):
    config_class = AetherMindConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["AetherMindDecoderLayer"]
    _supports_flash_attn_2 = False
    _supports_sdpa = True
    _supports_cache_class = True

    def _init_weights(self, module: nn.Module) -> None:
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


class AetherMindModel(AetherMindPreTrainedModel):
    """Transformer backbone without the LM head."""

    def __init__(self, config: AetherMindConfig) -> None:
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [AetherMindDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
        )
        self.norm = AetherMindRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.moe_aux_loss: Optional[torch.Tensor] = None
        self.post_init()

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

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

    # ------------------------------------------------------------------
    def _create_causal_mask(
        self,
        bsz: int,
        q_len: int,
        kv_len: int,
        device: torch.device,
        dtype: torch.dtype,
        attention_mask: Optional[torch.Tensor],
    ) -> torch.Tensor:
        """Additive (min-val) mask combining causality and padding."""
        min_val = torch.finfo(dtype).min
        q_pos = torch.arange(kv_len - q_len, kv_len, device=device).unsqueeze(-1)  # (q, 1)
        k_pos = torch.arange(0, kv_len, device=device).unsqueeze(0)                # (1, kv)
        causal = (k_pos > q_pos).unsqueeze(0).unsqueeze(0)                         # (1, 1, q, kv)
        mask = torch.zeros((bsz, 1, q_len, kv_len), dtype=dtype, device=device)
        mask = mask.masked_fill(causal, min_val)
        if attention_mask is not None:
            pad = (attention_mask[:, None, None, :kv_len] == 0)
            mask = mask.masked_fill(pad, min_val)
        return mask

    # ------------------------------------------------------------------
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[DynamicCache, List[Tuple[torch.Tensor, torch.Tensor]]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else False
        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else True
        if self.gradient_checkpointing and self.training:
            use_cache = False  # the cache is mutated in-place and breaks recomputation

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        bsz, seq_len, _ = inputs_embeds.shape
        device = inputs_embeds.device

        # ---- cache handling (accepts DynamicCache or legacy tuple list) ----
        if past_key_values is not None and not hasattr(past_key_values, "update"):
            legacy = past_key_values
            past_key_values = DynamicCache()
            for layer_idx, (k, v) in enumerate(legacy):
                past_key_values.update(k.to(device), v.to(device), layer_idx)
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
        if cache_position is None:
            cache_position = torch.arange(past_seen, past_seen + seq_len, device=device)
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # Explicit mask is required for padding, or for chunked prefill where the
        # fast is_causal=True kernel cannot express the offset (kv_len > q_len > 1).
        kv_len = past_seen + seq_len
        need_explicit_mask = attention_mask is not None or (seq_len > 1 and kv_len > seq_len)
        causal_mask = (
            self._create_causal_mask(bsz, seq_len, kv_len, device, inputs_embeds.dtype, attention_mask)
            if need_explicit_mask
            else None
        )

        hidden_states = inputs_embeds
        all_hidden_states: Tuple = ()
        all_self_attns: Tuple = ()
        aux_total: Optional[torch.Tensor] = None

        for decoder_layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    causal_mask,
                    position_ids,
                    past_key_values if use_cache else None,
                    use_cache,
                    cache_position,
                    output_attentions,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=causal_mask,
                    position_ids=position_ids,
                    past_key_values=past_key_values if use_cache else None,
                    use_cache=use_cache,
                    cache_position=cache_position,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]
            if output_attentions:
                all_self_attns += (layer_outputs[1],)
            if layer_outputs[2] is not None:
                aux_total = layer_outputs[2] if aux_total is None else aux_total + layer_outputs[2]

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        self.moe_aux_loss = aux_total

        if not return_dict:
            return tuple(
                v for v in (hidden_states, past_key_values, all_hidden_states, all_self_attns) if v is not None
            )
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states if output_hidden_states else None,
            attentions=all_self_attns if output_attentions else None,
        )


# ======================================================================
# Causal LM
# ======================================================================
class AetherMindForCausalLM(AetherMindPreTrainedModel, GenerationMixin):
    """AetherMind with a tied/untied LM head for causal language modeling."""

    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

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

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

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

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

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

    # ------------------------------------------------------------------
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[DynamicCache, List[Tuple[torch.Tensor, torch.Tensor]]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else False
        output_hidden_states = output_hidden_states if output_hidden_states is not None else False
        return_dict = return_dict if return_dict is not None else True

        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,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            cache_position=cache_position,
        )
        hidden_states = outputs.last_hidden_state

        if isinstance(logits_to_keep, int):
            slice_indices = slice(None, -logits_to_keep if logits_to_keep > 0 else None)
        else:
            slice_indices = logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :]).float()

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )
            if self.config.use_moe and self.model.moe_aux_loss is not None:
                loss = loss + self.config.moe_aux_loss_coeff * self.model.moe_aux_loss

        if not return_dict:
            output = (logits, outputs.past_key_values)
            if output_hidden_states:
                output += (outputs.hidden_states,)
            if output_attentions:
                output += (outputs.attentions,)
            return ((loss,) + output) if loss is not None else output

        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: Optional[DynamicCache] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        use_cache: bool = True,
        **kwargs,
    ):
        past_length = 0
        if past_key_values is not None and hasattr(past_key_values, "get_seq_length"):
            past_length = past_key_values.get_seq_length()

        model_inputs = {}
        if inputs_embeds is not None and past_length == 0:
            model_inputs["inputs_embeds"] = inputs_embeds
        elif cache_position is not None:
            model_inputs["input_ids"] = input_ids[:, cache_position].contiguous()
        else:
            model_inputs["input_ids"] = input_ids[:, past_length:].contiguous()
            cache_position = torch.arange(
                past_length, past_length + model_inputs["input_ids"].shape[1], device=input_ids.device
            )

        input_length = model_inputs["inputs_embeds"].shape[1] if "inputs_embeds" in model_inputs \
            else model_inputs["input_ids"].shape[1]

        position_ids = None
        if attention_mask is not None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            position_ids = position_ids[:, -input_length:]

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
                "position_ids": position_ids,
                "cache_position": cache_position,
            }
        )
        return model_inputs


__all__ = [
    "AetherMindConfig",
    "AetherMindModel",
    "AetherMindForCausalLM",
    "AetherMindPreTrainedModel",
    "AetherMindRMSNorm",
    "AetherMindRotaryEmbedding",
    "AetherMindAttention",
    "AetherMindMLP",
    "AetherMindSparseMoE",
    "AetherMindDecoderLayer",
]