Update model.py
e8a9009 verified about 5 hours ago
29.8 kB
# coding=utf-8
"""KenT-300M: mạng hồi quy một-token-một-bước với bộ nhớ EMA học được ở MỌI độ sâu.
Mỗi bước nhận ĐÚNG MỘT token x[t] (vector 1 x D); lịch sử nằm hoàn toàn trong trạng thái H (EMA). Sơ đồ:
    h_0 = x[t]                                   # embedding (Tokenizer)
    h_i = h_{i-1} + SwiGLU( RMSNorm(h_{i-1}) )   # i = 1..n   (khối xanh, residual)
    u_i = RMSNorm(h_i) · D^T · W_i               # i = 0..n   (khối cam "~DCT-II Proj (D^T)W": DCT-II cố định rồi W học được)
    H_i[t+1] = (1 - A_i) ⊙ H_i[t] + A_i ⊙ u_i    # "Learnable EMA Vector": A = sigmoid(ema_logit), shape (n+1, D)
    Hcat = Concat_i H_i[t+1]                     # (N, D), N = n + 1
    p    = DepthPool(Hcat)                       # RMSNorm -> Attn Pooling (query học được 1 x D) -> Query Residual, ra (1, D)
    g    = p + SwiGLU( RMSNorm(p) )              # khối xanh cuối
    y    = g · D                                 # "DCT-II Proj": DCT ngược (D trực chuẩn, D^T D = I)
    out  = x[t] + y                              # Add với đầu vào gốc x[t]
    logits = LM_head( RMSNorm_f(out) )           # head buộc chung trọng số với embedding
Attn Pooling theo độ sâu: mỗi hàng (tap) của Hcat được RMSNorm làm key/value;
query là MỘT vector học được (1, D) dùng chung mọi token, pooling toàn bộ N vector; "Query Residual" cộng ngược
query vào đầu ra attention.
Quy ước hàng: x là vector hàng nên DCT-II là x·D^T (D[k, i] = α_k cos(π(i+½)k/D), hàng = tần số k, cột = vị trí i)
và phép chiếu ngược là y·D. W khởi tạo ≈ ma trận đơn vị nên lúc đầu khối cam đúng bằng DCT-II và khối ra đúng
bằng DCT ngược.
Độ sâu (n khối xanh) cho N = n+1 "điểm lấy mẫu" (tap), mỗi tap có bộ nhớ EMA riêng H_i shape (D,). Trạng thái giải mã
toàn mô hình chỉ là H shape (N, D): O(1) theo độ dài chuỗi. D = độ phân giải (D lớn => chênh lệch góc/tần số nhỏ),
N = khả năng "tư duy" (N sâu => học quan hệ tốt hơn).
EMA là hồi quy tuyến tính theo từng kênh nê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ự. EMA và RMSNorm chạy ở FP32; trọng số lưu ở BF16.
RMSNorm cuối (`norm_f`) trước LM head không có trong sơ đồ; thêm vào để logits ổn định.
"""

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-300M.
    Args:
        vocab_size: kích thước từ vựng.
        hidden_size: D, độ rộng của vector token x[t] (cũng là độ rộng của mọi tap, của DCT và của EMA).
        num_hidden_layers: n, số khối xanh (SwiGLU residual) nối tiếp. Có n + 1 tap EMA (tap 0 lấy từ embedding).
        intermediate_size: kích thước hidden của mọi SwiGLU. Nên >= 2D để SwiGLU không thành nút cổ chai.
        pool_num_heads: số head của Attn Pooling theo độ sâu (hidden_size phải chia hết).
        rms_norm_eps: epsilon của RMSNorm.
        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 tap đầu, feature đầu (quên nhanh).
        ema_alpha_last: hệ số a khởi tạo của tap 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 chỉ số (tap, feature). a vẫn là tham số học được.
        use_cache: trả về KentCache khi suy luận.
        use_gradient_checkpointing: bật gradient checkpointing theo từng khối xanh khi khởi tạo model.
    """

    model_type = "kent"

    def __init__(
        self,
        vocab_size: int = 50257,
        hidden_size: int = 1152,
        num_hidden_layers: int = 24,
        intermediate_size: int = 2352,
        pool_num_heads: int = 9,
        rms_norm_eps: float = 1e-5,
        initializer_range: float = 0.02,
        ema_alpha_first: float = 0.5,
        ema_alpha_last: float = 0.01,
        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.intermediate_size = intermediate_size
        self.pool_num_heads = pool_num_heads
        self.rms_norm_eps = rms_norm_eps
        self.initializer_range = initializer_range
        self.ema_alpha_first = ema_alpha_first
        self.ema_alpha_last = ema_alpha_last
        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,
        )

    @property
    def num_taps(self) -> int:
        return self.num_hidden_layers + 1

    def alpha_init_matrix(self):
        """Ma trận a ((n+1) x D) khởi tạo: giảm theo cấp số nhân theo độ sâu của tap và theo chỉ số feature."""
        n_t, n_d = self.num_taps, self.hidden_size
        ratio = self.ema_alpha_last / self.ema_alpha_first
        return [[float(self.ema_alpha_first * ratio ** (0.5 * ((i / (n_t - 1) if n_t > 1 else 0.0)
                                                              + (j / (n_d - 1) if n_d > 1 else 0.0))))
                 for j in range(n_d)] for i in range(n_t)]

    def alpha_init(self, tap_idx: int) -> float:
        """Trung bình (theo feature) của a khởi tạo của tap `tap_idx` (để log)."""
        row = self.alpha_init_matrix()[tap_idx]
        return float(sum(row) / len(row))


# --------------------------------------------------------------------------- #
# Fixed DCT-II operator (trục cuối, vector hàng)
# --------------------------------------------------------------------------- #
@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 D (cố định, không phải tham số học): D[k, i] = α_k cos(π(i+½)k/n)."""
    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(x: torch.Tensor) -> torch.Tensor:
    """DCT-II của vector hàng trên trục cuối: x · D^T (tương đương D x cho vector cột)."""
    d = _dct2_matrix(x.shape[-1], x.device, x.dtype)
    return torch.matmul(x, d.t())


def idct2(x: torch.Tensor) -> torch.Tensor:
    """Phép chiếu ngược: x · D (D trực chuẩn nên idct2(dct2(x)) = x)."""
    d = _dct2_matrix(x.shape[-1], x.device, x.dtype)
    return torch.matmul(x, d)


# --------------------------------------------------------------------------- #
# EMA song song theo thời gian (affine-scan Hillis-Steele, hệ số riêng từng kênh)
# --------------------------------------------------------------------------- #
def _affine_scan(dec: torch.Tensor, acc: torch.Tensor):
    """Inclusive affine-scan Hillis-Steele theo trục 1: acc_t <- dec_t * acc_{t-1} + acc_t. Trả (acc, tích dec lũy tiến)."""
    t = acc.shape[1]
    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
    return acc, dec


class _EmaScanFn(torch.autograd.Function):
    """H_t = (1 - a m_t) ⊙ H_{t-1} + a m_t ⊙ u_t, backward GIẢI TÍCH (không lưu các tensor trung gian của scan log-depth).
    Autograd thường giữ ~2 log2(T) tensor (B,T,N,D) fp32 cho backward của scan (chiếm ~85% activation). Ở đây chỉ giữ u và H
    (2 tensor), backward là một scan ngược cùng cấu trúc:
        r_t = gH_t + (1 - a m_{t+1}) ⊙ r_{t+1}        (r = đạo hàm toàn phần theo H_t)
        dU_t = a m_t ⊙ r_t ;  dH0 = (1 - a m_0) ⊙ r_0 ;  dA = sum_{b,t} m_t r_t ⊙ (u_t - H_{t-1})
    """

    @staticmethod
    def forward(ctx, u, mask, h0, a):
        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, dec_cum = _affine_scan(dec, am * u)
        h = dec_cum * h0.unsqueeze(1) + acc
        ctx.save_for_backward(u, h, mask, h0, a)
        return h

    @staticmethod
    def backward(ctx, gh):
        u, h, mask, h0, a = ctx.saved_tensors
        bsz, t, n, d = u.shape
        m = mask.to(u.dtype).view(bsz, t, 1, 1)
        dec = 1.0 - a.view(1, 1, n, d) * m
        dec_f = dec.flip(1)
        coef = torch.cat([torch.zeros_like(dec_f[:, :1]), dec_f[:, :-1]], dim=1)   # hệ số của bước s (đã đảo trục) = dec_{t+1}
        r, _ = _affine_scan(coef, gh.flip(1))
        r = r.flip(1)                                                              # (B,T,N,D)
        du = a.view(1, 1, n, d) * m * r
        dh0 = dec[:, 0] * r[:, 0]
        h_prev = torch.cat([h0.unsqueeze(1), h[:, :-1]], dim=1)
        da = (m * r * (u - h_prev)).sum(dim=(0, 1))
        return du, None, dh0, da


def ema_scan(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: (N, D) khác nhau theo từng kênh; luôn chạy fp32.
    u: (B,T,N,D) | mask: (B,T) {0,1} (0 = padding: H giữ nguyên) | h0: (B,N,D).
    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 cho
    trạng thái sau mọi t. Mọi A_t nằm trong (0, 1] nên ổn định. Backward giải tích (`_EmaScanFn`) để tiết kiệm VRAM.
    Trả về H (B,T,N,D) và h_cuối (B,N,D).
    """
    with torch.autocast(device_type=u.device.type, enabled=False):
        h = _EmaScanFn.apply(u.float(), mask, h0.float(), a.float())
        return h, h[:, -1].clone()   # clone: cache không giữ sống cả tensor (B,T,N,D)


def ema_scan_reference(u: torch.Tensor, mask: torch.Tensor, h0: torch.Tensor, a: torch.Tensor):
    """Bản autograd thuần (chậm / tốn VRAM) của ema_scan, chỉ để kiểm tra gradient."""
    with torch.autocast(device_type=u.device.type, enabled=False):
        u, h0, a = u.float(), h0.float(), a.float()
        bsz, t, n, d = u.shape
        am = a.view(1, 1, n, d) * mask.to(u.dtype).view(bsz, t, 1, 1)
        acc, dec = _affine_scan(1.0 - am, am * u)
        h = dec * h0.unsqueeze(1) + acc
        return h, h[:, -1]


# --------------------------------------------------------------------------- #
# Cache (trạng thái H của toàn mô hình)
# --------------------------------------------------------------------------- #
class KentCache:
    """Trạng thái hồi quy của KenT: một tensor H (EMA, fp32) shape (B, n+1, D), giữ trong danh sách `h` (1 phần tử)."""

    def __init__(self, config: KentConfig, batch_size: int = 1, device=None):
        shape = (batch_size, config.num_taps, config.hidden_size)
        self.h = [torch.zeros(shape, dtype=torch.float32, device=device)]
        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: 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 KentBlock(nn.Module):
    """Khối xanh: h' = h + SwiGLU(RMSNorm(h)) (skip bắt đầu trước RMSNorm, như sơ đồ)."""

    def __init__(self, config: KentConfig):
        super().__init__()
        self.norm = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = KentSwiGLU(config.hidden_size, config.intermediate_size)

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


class KentTap(nn.Module):
    """Khối cam "~DCT-II Proj (D^T)W": u = RMSNorm(h) · D^T · W. D cố định, W (D x D) học được, khởi tạo ≈ I."""

    def __init__(self, config: KentConfig):
        super().__init__()
        self.norm = KentRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)  # W
        self.proj._identity_init = True

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        return self.proj(dct2(self.norm(h)))


class KentDepthPool(nn.Module):
    """Khối xanh dương: RMSNorm -> Attn Pooling -> Query Residual.
    Nhận ma trận EMA (B,T,N,D) (mỗi hàng một tap), trả (B,T,D). Mỗi hàng được RMSNorm làm key/value; query là một vector học được (1, D) dùng chung, pooling hết N vector.
    Đầu ra = query + Wo · attention (Query Residual).
    """

    def __init__(self, config: KentConfig):
        super().__init__()
        d = config.hidden_size
        if d % config.pool_num_heads != 0:
            raise ValueError("hidden_size phải chia hết cho pool_num_heads.")
        self.heads = config.pool_num_heads
        self.norm = KentRMSNorm(d, config.rms_norm_eps)
        self.query = nn.Parameter(torch.zeros(1, d))                    # learner query (1, D)
        self.q_proj = nn.Linear(d, d, bias=False)
        self.k_proj = nn.Linear(d, d, bias=False)
        self.v_proj = nn.Linear(d, d, bias=False)
        self.o_proj = nn.Linear(d, d, bias=False)
        self.o_proj._residual_out = True

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        b, t, n, d = h.shape
        x = self.norm(h)                                                # (B,T,N,D)
        query = self.query.to(h.dtype).expand(b, t, d)                  # (B,T,D): learner query dùng chung
        dh = d // self.heads
        q = self.q_proj(query).view(b, t, self.heads, dh)
        k = self.k_proj(x).view(b, t, n, self.heads, dh)
        v = self.v_proj(x).view(b, t, n, self.heads, dh)
        scores = torch.einsum("bthd,btnhd->bthn", q, k).float() * (dh ** -0.5)
        attn = torch.softmax(scores, dim=-1).to(v.dtype)
        out = torch.einsum("bthn,btnhd->bthd", attn, v).reshape(b, t, d)
        return query + self.o_proj(out)                                 # Query Residual


# --------------------------------------------------------------------------- #
# 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 = ["KentBlock"]
    _is_stateful = True

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        """Bật checkpoint theo khối xanh; mặc định use_reentrant=False."""
        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, KentDepthPool):
            nn.init.normal_(module.query, mean=0.0, std=std)
        elif isinstance(module, KentModel):
            with torch.no_grad():
                a0 = torch.tensor(self.config.alpha_init_matrix(), dtype=torch.float32)
                module.ema_logit.copy_(torch.log(a0 / (1.0 - a0)).to(module.ema_logit.device))
        if isinstance(module, nn.Linear) and getattr(module, "_identity_init", False):
            # W của khối cam ≈ I: khối cam khởi đầu đúng bằng DCT-II (làm ở mức Linear vì thứ tự duyệt module của HF)
            with torch.no_grad():
                w = module.weight
                w.copy_(torch.eye(w.shape[0], dtype=w.dtype, device=w.device) + std * torch.randn_like(w))
        if isinstance(module, nn.Linear) and getattr(module, "_residual_out", False):
            # đầu ra ghi vào residual stream => thu nhỏ theo (2 (n+1))^(-1/2) (kiểu GPT-2)
            with torch.no_grad():
                module.weight.mul_((2 * (self.config.num_hidden_layers + 1)) ** -0.5)


class KentModel(KentPreTrainedModel):
    def __init__(self, config: KentConfig):
        super().__init__(config)
        d = config.hidden_size
        self.embeddings = nn.Embedding(config.vocab_size, d)
        self.blocks = nn.ModuleList([KentBlock(config) for _ in range(config.num_hidden_layers)])
        self.taps = nn.ModuleList([KentTap(config) for _ in range(config.num_taps)])
        self.ema_logit = nn.Parameter(torch.zeros(config.num_taps, d))   # Learnable EMA Vector: a = sigmoid(logit)
        self.pool = KentDepthPool(config)                                # RMSNorm -> Attn Pooling -> Query Residual
        self.out_block = KentBlock(config)                               # khối xanh cuối (trước DCT-II Proj ngược)
        self.norm_f = KentRMSNorm(d, config.rms_norm_eps)
        for blk in self.blocks:
            blk.mlp.down._residual_out = True
        self.out_block.mlp.down._residual_out = True
        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 alpha(self) -> torch.Tensor:
        """Learnable EMA a ((n+1) x D) trong (0,1), tính ở fp32."""
        return torch.sigmoid(self.ema_logit.float()).clamp_max(0.999)

    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.
        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, d = 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)

        # --- chuỗi khối xanh (mọi token độc lập nhau) -> n+1 điểm lấy mẫu h_0..h_n
        all_hidden = (x,) if output_hidden_states else None
        feats = [x]
        for blk in self.blocks:
            if self.gradient_checkpointing and self.training:
                x = self._gradient_checkpointing_func(blk.__call__, x)
            else:
                x = blk(x)
            feats.append(x)
            if output_hidden_states:
                all_hidden += (x,)

        # --- khối cam: u_i = RMSNorm(h_i) D^T W_i, rồi EMA riêng từng tap
        u = torch.stack([tap(h) for tap, h in zip(self.taps, feats)], dim=2).float()   # (B,T,n+1,D)
        if chunked:  # K chunk liên tiếp của một tài liệu: nối thành một chuỗi cho EMA
            u = u.reshape(1, b * t, self.config.num_taps, d)
            m = mask.reshape(1, b * t)
        else:
            m = mask
        h0 = cache_params.h[0] if cache_params is not None else u.new_zeros((u.shape[0], self.config.num_taps, d))
        h, h_last = ema_scan(u, m, h0, self.alpha())                                    # (.,.,n+1,D), (.,n+1,D)
        if use_cache:
            cache_params.h[0] = h_last
            cache_params.steps += b * t if chunked else t

        # --- Concat [N, D] -> RMSNorm -> Attn Pooling -> Query Residual -> khối xanh -> DCT-II Proj -> Add x[t]
        p = self.pool(h.to(feats[0].dtype))                                             # (.,.,D)
        p = self.out_block(p)
        y = idct2(p)
        x0 = feats[0].reshape(1, b * t, d) if chunked else feats[0]                     # đầu vào gốc x[t] (embedding)
        y = self.norm_f(x0 + y)                                                         # (1, K*C, D) nếu chunked
        if output_hidden_states:
            all_hidden += (y,)

        if not return_dict:
            return tuple(v for v in (y, cache_params if use_cache else None, all_hidden) if v is not None)
        return KentOutput(
            last_hidden_state=y,
            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",
    "idct2",
    "ema_scan",
    "ema_scan_reference",
]