# coding=utf-8
"""KenT: kiến trúc hồi quy hai luồng (representation X_t / temporal-memory H_t).

Mỗi bước nhận ĐÚNG MỘT token (cửa sổ trượt từng token một); lịch sử nằm trong trạng thái H của từng layer.
Mỗi KenT layer (hộp "WNN" trong sơ đồ) thực hiện:

    Z    = ResProj_in(Norm_X(X))                # 1 -> M: một vector token thành M resolution, mỗi resolution d feature
    Z    = Z + SwiGLU( DCT_M( RMSNorm(Z) ) )    # WNN Block: DCT-II cố định trên trục M, SwiGLU trên trục d, skip quanh cả cụm
    H_t  = (1 - A) ⊙ H_{t-1} + A ⊙ Z_t          # "EMA Learner Matrix": A = sigmoid(ema_logit), MA TRẬN M x d học được / layer
    X'   = X + ResProj_out( DCT_M( Norm_H(H_t) ) )   # DCT-II Proj rồi M -> 1; Add với X vào layer (residual)

Mỗi layer xếp K = `num_wnn_blocks` WNN Block nối tiếp (mặc định K = 1, đúng sơ đồ); độ sâu đến từ việc xếp chồng
nhiều layer ("xN" trong sơ đồ). Layer cuối cho X' -> norm cuối -> lm_head (tied với embedding).

Bộ nhớ H có cơ chế quên: A là ma trận M x d (mỗi vị trí resolution VÀ mỗi feature có hệ số quên riêng), học được và
riêng từng layer. Khởi tạo: nội suy theo cấp số nhân từ `ema_alpha_first` (layer đầu, vị trí đầu, feature đầu) tới
`ema_alpha_last` (layer cuối, vị trí cuối, feature cuối) theo trung bình của chỉ số layer, vị trí và feature.
`ema_per_feature=False` quay về vector M phần tử (chung cho mọi feature), dùng scan theo khối.

EMA là một phép hồi quy tuyến tính (đường chéo theo từng kênh) nên vẫn song song hóa được theo thời gian bằng
affine-scan log-depth (Hillis-Steele): kết quả bằng đúng hồi quy tuần tự. Trạng thái giải mã chỉ là H cho mỗi
layer, shape (M, d): O(1) theo độ dài.

Trọng số được lưu ở FP16; EMA và normalization chạy ở FP32.
"""

import math
from dataclasses import dataclass
from functools import lru_cache
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import CrossEntropyLoss

from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
from transformers.utils import ModelOutput


# --------------------------------------------------------------------------- #
# Config
# --------------------------------------------------------------------------- #
class KentConfig(PretrainedConfig):
    r"""Cấu hình KenT.

    Args:
        vocab_size: kích thước từ vựng.
        hidden_size: số chiều của representation stream X_t (n trong sơ đồ).
        num_hidden_layers: số KenT layer xếp chồng.
        resolution_size: M, độ dài trục mà DCT-II tác động lên (cũng là độ dài vector hệ số EMA, "d=M" trong sơ đồ).
        feature_dim: d, số chiều đặc trưng trên mỗi vị trí resolution.
        num_wnn_blocks: K, số WNN Block (DCT -> SwiGLU -> skip) nối tiếp trong mỗi layer. Sơ đồ: K = 1.
        intermediate_size: kích thước hidden của SwiGLU (trên trục d). Nên >= 2d để SwiGLU không thành nút cổ chai.
        rms_norm_eps: epsilon của RMSNorm.
        use_bias: dùng bias trong Resolution Proj (SwiGLU không bias).
        initializer_range: độ lệch chuẩn khởi tạo trọng số.
        ema_alpha_first: hệ số a khởi tạo của layer đầu, vị trí đầu (quên nhanh).
        ema_alpha_last: hệ số a khởi tạo của layer cuối, vị trí cuối, feature cuối (nhớ lâu); phần còn lại nội suy theo
            cấp số nhân theo trung bình các chỉ số (layer, vị trí[, feature]). a vẫn là tham số học được.
        ema_per_feature: True = EMA Learner Matrix (hệ số quên riêng cho mỗi (vị trí M, feature d), ema_logit shape
            (M, d)); False = vector M phần tử, chung cho mọi feature (ema_logit shape (M,)).
        ema_block_size: độ dài khối của EMA song song (chỉ dùng khi ema_per_feature=False; ảnh hưởng tốc độ/bộ nhớ).
        use_cache: trả về KentCache khi suy luận.
        use_gradient_checkpointing: bật gradient checkpointing theo từng KenT layer khi khởi tạo model
            (chỉ có tác dụng ở chế độ train; tương đương gọi `model.gradient_checkpointing_enable()`).
    """

    model_type = "kent"

    def __init__(
        self,
        vocab_size: int = 50257,
        hidden_size: int = 256,
        num_hidden_layers: int = 24,
        resolution_size: int = 12,
        feature_dim: int = 96,
        num_wnn_blocks: int = 1,
        intermediate_size: int = 384,
        rms_norm_eps: float = 1e-5,
        use_bias: bool = True,
        initializer_range: float = 0.02,
        ema_alpha_first: float = 0.5,
        ema_alpha_last: float = 0.01,
        ema_per_feature: bool = True,
        ema_block_size: int = 128,
        use_cache: bool = True,
        use_gradient_checkpointing: bool = False,
        pad_token_id: int = 50256,
        bos_token_id: int = 50256,
        eos_token_id: int = 50256,
        tie_word_embeddings: bool = True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.resolution_size = resolution_size
        self.feature_dim = feature_dim
        self.num_wnn_blocks = num_wnn_blocks
        self.intermediate_size = intermediate_size
        self.rms_norm_eps = rms_norm_eps
        self.use_bias = use_bias
        self.initializer_range = initializer_range
        self.ema_alpha_first = ema_alpha_first
        self.ema_alpha_last = ema_alpha_last
        self.ema_per_feature = ema_per_feature
        self.ema_block_size = ema_block_size
        self.use_cache = use_cache
        self.use_gradient_checkpointing = use_gradient_checkpointing
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )

    def alpha_init_vector(self, layer_idx: int):
        """Vector a (M phần tử) khởi tạo của layer `layer_idx`: giảm theo cấp số nhân theo độ sâu và theo vị trí."""
        n_l, n_p = self.num_hidden_layers, self.resolution_size
        r_l = layer_idx / (n_l - 1) if n_l > 1 else 0.0
        ratio = self.ema_alpha_last / self.ema_alpha_first
        return [float(self.ema_alpha_first * ratio ** (0.5 * (r_l + (p / (n_p - 1) if n_p > 1 else 0.0))))
                for p in range(n_p)]

    def alpha_init_matrix(self, layer_idx: int):
        """Ma trận a (M x d) khởi tạo của layer: giảm theo cấp số nhân theo độ sâu, vị trí và feature."""
        n_l, n_p, n_d = self.num_hidden_layers, self.resolution_size, self.feature_dim
        r_l = layer_idx / (n_l - 1) if n_l > 1 else 0.0
        ratio = self.ema_alpha_last / self.ema_alpha_first
        return [[float(self.ema_alpha_first * ratio ** ((r_l + (p / (n_p - 1) if n_p > 1 else 0.0)
                                                         + (j / (n_d - 1) if n_d > 1 else 0.0)) / 3.0))
                 for j in range(n_d)] for p in range(n_p)]

    def alpha_init(self, layer_idx: int) -> float:
        """Trung bình của a khởi tạo của layer (để log)."""
        if self.ema_per_feature:
            flat = [x for row in self.alpha_init_matrix(layer_idx) for x in row]
        else:
            flat = self.alpha_init_vector(layer_idx)
        return float(sum(flat) / len(flat))


# --------------------------------------------------------------------------- #
# Fixed DCT-II operator
# --------------------------------------------------------------------------- #
@lru_cache(maxsize=32)
def _dct2_matrix(n: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Ma trận DCT-II trực chuẩn A_N (cố định, không phải tham số học)."""
    k = torch.arange(n, dtype=torch.float64).unsqueeze(1)
    i = torch.arange(n, dtype=torch.float64).unsqueeze(0)
    a = torch.cos(math.pi / n * (i + 0.5) * k)
    alpha = torch.full((n, 1), math.sqrt(2.0 / n), dtype=torch.float64)
    alpha[0, 0] = math.sqrt(1.0 / n)
    return (alpha * a).to(device=device, dtype=dtype)


def dct2(z: torch.Tensor) -> torch.Tensor:
    """D_N(Z) = A_N Z, tác động lên trục -2 (độ dài N) của tensor (..., N, d_h)."""
    a = _dct2_matrix(z.shape[-2], z.device, z.dtype)
    return torch.matmul(a, z)


# --------------------------------------------------------------------------- #
# EMA: song song theo thời gian (khối tam giác + affine-scan giữa các khối)
# --------------------------------------------------------------------------- #
def _affine_scan(a: torch.Tensor, b: torch.Tensor):
    """Inclusive scan dọc dim 1 của các ánh xạ affine h -> a*h + b (Hillis-Steele, log K bước).

    a: (B,K,N) | b: (B,K,N,D). Phần tử k được áp dụng SAU các phần tử < k.
    Trả về (A, B) sao cho trạng thái sau khối k là A[k]*h_vào + B[k].
    """
    k = a.shape[1]
    off = 1
    while off < k:
        b = torch.cat([b[:, :off], a[:, off:, :, None] * b[:, :-off] + b[:, off:]], dim=1)
        a = torch.cat([a[:, :off], a[:, off:] * a[:, :-off]], dim=1)
        off *= 2
    return a, b


def ema_scan(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
    """EMA song song, luôn chạy fp32 (tắt autocast để matmul không bị ép về fp16).

    a: (M, d) -> hệ số riêng từng kênh (EMA Learner Matrix, scan Hillis-Steele); a: (M,) -> theo khối (chung cho d).
    """
    with torch.autocast(device_type=u.device.type, enabled=False):
        if a.dim() == 2:
            return _ema_scan_diag(u.float(), mask, h0.float(), a.float())
        return _ema_scan_fp32(u.float(), mask, h0.float(), a.float(), block)


def _ema_scan_diag(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor):
    """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t với a: (M, d) khác nhau theo từng kênh.

    u: (B,T,M,d) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,M,d) fp32.
    Mỗi bước là ánh xạ affine h -> A_t h + B_t (A_t = 1 - a m_t, B_t = a m_t u_t); inclusive scan log-depth
    (Hillis-Steele) cho trạng thái sau mọi t. Mọi A_t nằm trong (0, 1] nên ổn định.
    Trả về H (B,T,M,d) và h_cuối (B,M,d).
    """
    bsz, t, n, d = u.shape
    am = a.view(1, 1, n, d) * mask.to(u.dtype).view(bsz, t, 1, 1)
    dec = 1.0 - am
    acc = am * u
    off = 1
    while off < t:
        acc = torch.cat([acc[:, :off], dec[:, off:] * acc[:, :-off] + acc[:, off:]], dim=1)
        dec = torch.cat([dec[:, :off], dec[:, off:] * dec[:, :-off]], dim=1)
        off *= 2
    h = dec * h0.unsqueeze(1) + acc
    return h, h[:, -1]


def _ema_scan_fp32(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor, block: int = 128):
    """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t, tính song song cho mọi t.

    u: (B,T,N,D) fp32 | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D) fp32
    a: (N,) fp32 trong (0,1): hệ số riêng cho từng vị trí N (= M trong sơ đồ), chung cho mọi feature D.
    Trả về H (B,T,N,D) và h_cuối (B,N,D).

    Trong một khối (với từng vị trí n):
        H_t = sum_{s<=t} a m_s [prod_{s<r<=t}(1 - a m_r)] u_s + [prod_{r<=t}(1 - a m_r)] h_vào,
    với tích tính qua hiệu của cumsum log (luôn <= 0 nên ổn định). Các khối được nối bằng affine-scan.
    """
    bsz, t, n, d = u.shape
    blk = max(1, min(block, t))
    pad = (-t) % blk
    m = mask.to(u.dtype)
    if pad:  # padding ở cuối: mask = 0 nên trạng thái không đổi
        u = F.pad(u, (0, 0, 0, 0, 0, pad))
        m = F.pad(m, (0, pad))
    k = (t + pad) // blk
    uk = u.reshape(bsz, k, blk, n, d).permute(0, 1, 3, 2, 4)                  # (B,K,N,blk,D)
    am = a.view(1, 1, n, 1) * m.reshape(bsz, k, 1, blk)                       # (B,K,N,blk)

    cum = torch.cumsum(torch.log1p(-am), dim=-1)                              # log prod_{r<=t}(1 - a m_r)
    tril = torch.ones(blk, blk, dtype=torch.bool, device=u.device).tril()
    decay = torch.exp((cum.unsqueeze(-1) - cum.unsqueeze(-2)).masked_fill(~tril, float("-inf")))  # (B,K,N,blk,blk)
    local = torch.matmul(decay * am.unsqueeze(-2), uk)                        # (B,K,N,blk,D): khối bắt đầu từ H = 0

    a_blk = torch.exp(cum[..., -1])                                           # (B,K,N): hệ số phân rã cả khối
    a_p, b_p = _affine_scan(a_blk, local[:, :, :, -1])                        # b: (B,K,N,D)
    s = a_p.unsqueeze(-1) * h0.unsqueeze(1) + b_p                             # trạng thái sau mỗi khối (B,K,N,D)
    h_in = torch.cat([h0.unsqueeze(1), s[:, :-1]], dim=1)                     # trạng thái vào mỗi khối
    h = local + torch.exp(cum).unsqueeze(-1) * h_in.unsqueeze(3)              # (B,K,N,blk,D)
    h = h.permute(0, 1, 3, 2, 4).reshape(bsz, k * blk, n, d)[:, :t]
    return h, s[:, -1]


# --------------------------------------------------------------------------- #
# Cache (trạng thái H cho từng layer)
# --------------------------------------------------------------------------- #
class KentCache:
    """Trạng thái hồi quy của KenT: H (EMA, fp32) theo layer."""

    def __init__(self, config: KentConfig, batch_size: int = 1, device=None):
        shape = (batch_size, config.resolution_size, config.feature_dim)
        self.h = [
            torch.zeros(shape, dtype=torch.float32, device=device)
            for _ in range(config.num_hidden_layers)
        ]
        self.steps = 0  # số vị trí đã nạp (đếm ở phía Python, không cần sync GPU)

    @property
    def batch_size(self) -> int:
        return self.h[0].shape[0]

    @property
    def seen_tokens(self) -> int:
        return self.steps

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self.steps

    def reset(self):
        for h in self.h:
            h.zero_()
        self.steps = 0

    def clone(self) -> "KentCache":
        new = KentCache.__new__(KentCache)
        new.h = [h.clone() for h in self.h]
        new.steps = self.steps
        return new

    def to(self, device):
        self.h = [h.to(device) for h in self.h]
        return self

    def reorder_cache(self, beam_idx: torch.LongTensor):
        self.h = [h.index_select(0, beam_idx.to(h.device)) for h in self.h]


# --------------------------------------------------------------------------- #
# Building blocks
# --------------------------------------------------------------------------- #
class KentRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.weight.float()).to(dtype)


class KentSwiGLU(nn.Module):
    """SwiGLU trên trục d: down( silu(x W_g) ⊙ (x W_u) ), không bias."""

    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        g, u = self.gate_up(x).chunk(2, dim=-1)
        return self.down(F.silu(g) * u)


class KentWNNBlock(nn.Module):
    """WNN Block: Z' = Z + SwiGLU( DCT_M( RMSNorm(Z) ) ).

    RMSNorm theo trục d; DCT-II cố định tác động lên trục M ("Transpose DCT-II Proj": chuyển trục rồi chiếu);
    SwiGLU tác động lên trục d (dùng chung cho mọi vị trí M); skip nối quanh cả cụm Norm -> DCT -> SwiGLU.
    """

    def __init__(self, config: KentConfig):
        super().__init__()
        hidden = config.intermediate_size or 4 * config.feature_dim
        self.norm = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
        self.mlp = KentSwiGLU(config.feature_dim, hidden)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return z + self.mlp(dct2(self.norm(z)))


class KentLayer(nn.Module):
    """Một bước KenT: (X_t, H_{t-1}) -> (X_{t+1}, H_t), với K WNN Block và EMA có vector hệ số M phần tử."""

    def __init__(self, config: KentConfig, layer_idx: int = 0):
        super().__init__()
        self.n = config.resolution_size       # M
        self.d = config.feature_dim
        self.block_size = config.ema_block_size
        self.norm_x = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.res_in = nn.Linear(config.hidden_size, self.n * self.d, bias=config.use_bias)    # Resolution Proj n -> M
        self.wnn = nn.ModuleList([KentWNNBlock(config) for _ in range(config.num_wnn_blocks)])
        self.norm_h = KentRMSNorm(config.feature_dim, config.rms_norm_eps)
        self.res_out = nn.Linear(self.n * self.d, config.hidden_size, bias=config.use_bias)   # Resolution Proj M -> 1
        a0 = torch.tensor(
            config.alpha_init_matrix(layer_idx) if config.ema_per_feature else config.alpha_init_vector(layer_idx),
            dtype=torch.float32)
        self.ema_logit_init = torch.log(a0 / (1.0 - a0))
        self.ema_logit = nn.Parameter(self.ema_logit_init.clone())             # (M, d) hoặc (M,)

    def alpha(self) -> torch.Tensor:
        """EMA Learner Matrix a (M, d) (hoặc vector (M,)) trong (0,1), tính ở fp32."""
        return torch.sigmoid(self.ema_logit.float()).clamp_max(0.999)

    def forward(self, x, mask, h0, chunked: bool = False):
        b, t, _ = x.shape
        z = self.res_in(self.norm_x(x)).view(b, t, self.n, self.d)
        for blk in self.wnn:                          # K WNN Block nối tiếp
            z = blk(z)
        u = z.float()                                 # EMA ở fp32
        if chunked:  # trục batch là các chunk liên tiếp: nối thành một chuỗi cho EMA
            u = u.reshape(1, b * t, self.n, self.d)
            m = mask.reshape(1, b * t)
        else:
            m = mask
        if h0 is None:
            h0 = u.new_zeros((u.shape[0], self.n, self.d))
        h, h_last = ema_scan(u, m, h0, self.alpha(), self.block_size)
        h = h.reshape(b, t, self.n, self.d)
        spec = dct2(self.norm_h(h).to(x.dtype))       # DCT-II Proj trên trục M của H
        dx = self.res_out(spec.reshape(b, t, self.n * self.d))
        return x + dx, h_last


# --------------------------------------------------------------------------- #
# Outputs
# --------------------------------------------------------------------------- #
@dataclass
class KentOutput(ModelOutput):
    last_hidden_state: Optional[torch.FloatTensor] = None
    cache_params: Optional[KentCache] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None


@dataclass
class KentCausalLMOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    cache_params: Optional[KentCache] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None


# --------------------------------------------------------------------------- #
# Models
# --------------------------------------------------------------------------- #
class KentPreTrainedModel(PreTrainedModel):
    config_class = KentConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["KentLayer"]
    _is_stateful = True

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        """Bật checkpoint theo layer; mặc định use_reentrant=False (an toàn với args không phải tensor
        như h0=None / chunked, và không cần input requires_grad)."""
        if gradient_checkpointing_kwargs is None:
            gradient_checkpointing_kwargs = {"use_reentrant": False}
        super().gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, KentRMSNorm):
            nn.init.ones_(module.weight)
        elif isinstance(module, KentLayer):
            with torch.no_grad():
                module.ema_logit.copy_(module.ema_logit_init.to(module.ema_logit.device))
        if isinstance(module, nn.Linear) and getattr(module, "_residual_out", False):
            # đầu ra ghi vào residual stream => thu nhỏ theo (2L)^(-1/2) (kiểu GPT-2)
            with torch.no_grad():
                module.weight.mul_((2 * self.config.num_hidden_layers) ** -0.5)


class KentModel(KentPreTrainedModel):
    def __init__(self, config: KentConfig):
        super().__init__(config)
        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([KentLayer(config, i) for i in range(config.num_hidden_layers)])
        for layer in self.layers:
            layer.res_out._residual_out = True
        self.norm_f = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.post_init()
        if getattr(config, "use_gradient_checkpointing", False):
            self.gradient_checkpointing_enable()

    def get_input_embeddings(self):
        return self.embeddings

    def set_input_embeddings(self, value):
        self.embeddings = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        cache_params: Optional[KentCache] = None,
        use_cache: Optional[bool] = None,
        chunked: bool = False,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, KentOutput]:
        """
        chunked=False: trục batch là các chuỗi độc lập; (B, T).
        chunked=True : trục batch là K chunk liên tiếp của MỘT tài liệu; (K, C).
                       Cache vào/ra có batch size 1 (trạng thái sau chunk cuối).
        attention_mask: 1 = token thật, 0 = padding (H giữ nguyên tại vị trí đó).
        """
        use_cache = use_cache if use_cache is not None else (self.config.use_cache and not self.training)
        output_hidden_states = bool(output_hidden_states)
        return_dict = return_dict if return_dict is not None else True

        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("Cần đúng một trong input_ids hoặc inputs_embeds.")
        x = self.embeddings(input_ids) if inputs_embeds is None else inputs_embeds
        b, t, _ = x.shape

        if attention_mask is None:
            mask = torch.ones(b, t, dtype=torch.long, device=x.device)
        else:
            mask = attention_mask.to(device=x.device, dtype=torch.long)
            if mask.shape[1] != t:  # generate() truyền mask toàn chuỗi
                mask = mask[:, -t:]

        if cache_params is not None and cache_params.batch_size != (1 if chunked else b):
            raise ValueError("Batch size của cache không khớp với đầu vào.")
        if use_cache and cache_params is None:
            cache_params = KentCache(self.config, 1 if chunked else b, x.device)

        all_hidden = () if output_hidden_states else None
        for i, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden += (x,)
            h0 = cache_params.h[i] if cache_params is not None else None
            if self.gradient_checkpointing and self.training:
                x, h_last = self._gradient_checkpointing_func(layer.__call__, x, mask, h0, chunked)
            else:
                x, h_last = layer(x, mask, h0, chunked)
            if use_cache:
                cache_params.h[i] = h_last
        if use_cache:
            cache_params.steps += b * t if chunked else t

        x = self.norm_f(x)
        if chunked:  # nối các chunk lại thành một chuỗi (1, K*C, d)
            x = x.reshape(1, b * t, -1)
            if all_hidden is not None:
                all_hidden = tuple(h.reshape(1, b * t, -1) for h in all_hidden)
        if all_hidden is not None:
            all_hidden += (x,)

        if not return_dict:
            return tuple(v for v in (x, cache_params, all_hidden) if v is not None)
        return KentOutput(
            last_hidden_state=x,
            cache_params=cache_params if use_cache else None,
            hidden_states=all_hidden,
        )


class KentForCausalLM(KentPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embeddings.weight"}

    def __init__(self, config: KentConfig):
        super().__init__(config)
        self.model = KentModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def new_cache(self, batch_size: int = 1) -> KentCache:
        p = next(self.parameters())
        return KentCache(self.config, batch_size, p.device)

    # ---- tích hợp generate() ------------------------------------------------
    def prepare_inputs_for_generation(
        self,
        input_ids,
        inputs_embeds=None,
        attention_mask=None,
        use_cache=None,
        cache_params=None,
        past_key_values=None,
        **kwargs,
    ):
        if cache_params is None and isinstance(past_key_values, KentCache):
            cache_params = past_key_values
        if cache_params is not None and cache_params.steps > 0:
            input_ids = input_ids[:, -1:]  # trạng thái đã chứa phần còn lại
            inputs_embeds = None
        inputs = {"inputs_embeds": inputs_embeds} if inputs_embeds is not None else {"input_ids": input_ids}
        inputs.update(
            attention_mask=attention_mask,
            cache_params=cache_params,
            use_cache=True if use_cache is None else use_cache,
        )
        return inputs

    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False, **kwargs):
        model_kwargs["cache_params"] = outputs.get("cache_params", None)
        model_kwargs.pop("past_key_values", None)
        if model_kwargs.get("attention_mask", None) is not None:
            am = model_kwargs["attention_mask"]
            model_kwargs["attention_mask"] = torch.cat([am, am.new_ones((am.shape[0], 1))], dim=-1)
        if model_kwargs.get("cache_position", None) is not None:
            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1
        return model_kwargs

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        cache_params: Optional[KentCache] = None,
        past_key_values: Optional[KentCache] = None,
        use_cache: Optional[bool] = None,
        chunked: bool = False,
        labels: Optional[torch.LongTensor] = None,
        logits_to_keep: int = 0,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, KentCausalLMOutput]:
        if cache_params is None and isinstance(past_key_values, KentCache):
            cache_params = past_key_values
        return_dict = return_dict if return_dict is not None else True
        out = self.model(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_params=cache_params,
            use_cache=use_cache,
            chunked=chunked,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        hidden = out.last_hidden_state
        if labels is None and logits_to_keep:
            hidden = hidden[:, -logits_to_keep:]
        logits = self.lm_head(hidden)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous().float()
            shift_labels = labels[..., 1:].contiguous().to(shift_logits.device)
            loss = CrossEntropyLoss()(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        if not return_dict:
            return tuple(v for v in (loss, logits, out.cache_params, out.hidden_states) if v is not None)
        return KentCausalLMOutput(
            loss=loss,
            logits=logits,
            cache_params=out.cache_params,
            hidden_states=out.hidden_states,
        )


__all__ = [
    "KentConfig",
    "KentCache",
    "KentModel",
    "KentForCausalLM",
    "KentPreTrainedModel",
    "dct2",
    "ema_scan",
]