Blame for patch.py

"""transformers remote code for the exp_train hybrid LM checkpoints (trust_remote_code=True).

    tok = AutoTokenizer.from_pretrained(path)
    model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda")
    out = model.generate(**tok("The capital of France is", return_tensors="pt").to("cuda"), max_new_tokens=5)

Architecture (exp_train / a100stack): 32 layers, 24 Gated DeltaNet + 8 gated GQA softmax attention (every 4th),
SwiGLU MLPs, gated residuals, Delta Block attention residuals ("repo" variant) and a hashed bigram/trigram memory
(ngram.safetensors, canonical-token hashing, +-1 sign trick) injected before blocks 2 and 16.

Parameter names are the training checkpoint's, so model.safetensors loads unchanged; the n-gram table is read from
ngram.safetensors after loading. Matrices run in BF16, norms / gates / the residual stream in FP32 as in training.
Inference only, batch size 1. Prefill uses FLA's chunked gated-delta-rule kernel; decoding uses the fused recurrent
kernel, a preallocated KV cache, and a CUDA graph for the single-token step.
"""
import hashlib
import json
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast


class HybridLMConfig(PretrainedConfig):
    model_type = "exp_train_hybrid"

    def __init__(self, model_config=None, modifications=None, kv_cache_len=4096, use_cuda_graph=True,
                 bos_token_id=2, eos_token_id=3, pad_token_id=0, tie_word_embeddings=False, **kwargs):
        self.model_config = model_config or {}
        self.modifications = modifications or {}
        self.kv_cache_len = kv_cache_len          # KV-cache capacity (prompt + generated tokens)
        self.use_cuda_graph = use_cuda_graph        # capture the single-token decode step in a CUDA graph
        mc = self.model_config
        if mc:
            self.vocab_size = mc["vocab_size"]
            self.hidden_size = mc["d_model"]
            self.num_hidden_layers = mc["n_layers"]
            self.num_attention_heads = mc["n_heads"]
            self.num_key_value_heads = mc["n_kv_heads"]
            self.head_dim = mc["head_dim"]
            self.intermediate_size = mc["ffn_dim"]
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id,
                         tie_word_embeddings=tie_word_embeddings, **kwargs)


# ----------------------------------------------------------------------------- n-gram hashing (ngram.py)
_MULTS = [0x9E3779B97F4A7C15, 0xC2B2AE3D27D4EB4F, 0x165667B19E3779F9, 0xD6E8FEB86659FD93,
          0xFF51AFD7ED558CCD, 0xC4CEB9FE1A85EC53, 0x94D049BB133111EB, 0xBF58476D1CE4E5B9]
_SIGN_MULTS = [0xA0761D6478BD642F, 0xE7037ED1A0B428DB, 0x8EBC6AF09C88C6E3, 0x589965CC75374CC3,
               0x1D8E4E27C47D124F, 0x2D358DCCAA6C78A5, 0x9FB21C651E98DF25, 0xD3833E804F4C574B]


def _i64(x):
    return x - (1 << 64) if x >= (1 << 63) else x


def _shift(ids, j):   # tok[t] = ids[t - j]
    return torch.cat([ids.new_zeros(j), ids[: ids.numel() - j]]) if j else ids


def build_canon_map(tokenizer_file):
    """ngram_canon.py: token ids whose text normalizes to the same string share n-gram rows (hashing only)."""
    from tokenizers import Regex, Tokenizer, normalizers
    tokenizer = Tokenizer.from_file(str(tokenizer_file))
    norm = normalizers.Sequence([
        normalizers.NFKC(), normalizers.NFD(), normalizers.StripAccents(), normalizers.Lowercase(),
        normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), normalizers.Replace(Regex(r"^ $"), ""),
        normalizers.Strip(), normalizers.Replace("", " ")])
    vocab = tokenizer.get_vocab_size(with_added_tokens=True)
    special = {t["id"] for t in json.loads(tokenizer.to_str()).get("added_tokens", []) if t.get("special")}
    first, canon = {}, torch.arange(vocab, dtype=torch.int64)
    for t in range(vocab):
        if t in special:
            continue
        text = tokenizer.decode([t], skip_special_tokens=False)
        if "�" in text:
            continue
        key = norm.normalize_str(text)
        if key:
            canon[t] = first.setdefault(key, t)
    return canon


def _rms(x, w, eps):
    return F.rms_norm(x.float(), (x.shape[-1],), w.float(), eps)


def _select(sources, qw, eps):
    """Delta Block: sum_i softmax_i(q . RMSNorm_w(s_i)) s_i  (qw = query * weight)."""
    logits = torch.stack([(s * qw).sum(-1) * torch.rsqrt(s.pow(2).mean(-1) + eps) for s in sources])
    p = logits.softmax(0)
    out = p[0].unsqueeze(-1) * sources[0]
    for pi, s in zip(p[1:], sources[1:]):
        out = out + pi.unsqueeze(-1) * s
    return out


# ----------------------------------------------------------------------------- modules (checkpoint names)
def _linear(i, o):
    return nn.Linear(i, o, bias=False)


class _Weight(nn.Module):
    def __init__(self, *shape):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(*shape))


class _Gate(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.a = nn.Parameter(torch.zeros(d))


class GatedAttention(nn.Module):
    def __init__(self, c):
        super().__init__()
        H, KV, hd, d = c["n_heads"], c["n_kv_heads"], c["head_dim"], c["d_model"]
        self.q_proj, self.gate_proj, self.o_proj = _linear(d, H * hd), _linear(d, H * hd), _linear(H * hd, d)
        self.k_proj, self.v_proj = _linear(d, KV * hd), _linear(d, KV * hd)
        self.q_norm, self.k_norm = _Weight(hd), _Weight(hd)


class GatedDeltaNet(nn.Module):
    def __init__(self, c):
        super().__init__()
        H, dk, d = c["gdn_heads"], c["gdn_head_dim"], c["d_model"]
        dv, K = int(dk * c["gdn_expand_v"]), c["gdn_conv_size"]
        self.q_proj, self.k_proj = _linear(d, H * dk), _linear(d, H * dk)
        self.v_proj, self.g_proj, self.o_proj = _linear(d, H * dv), _linear(d, H * dv), _linear(H * dv, d)
        self.a_proj, self.b_proj = _linear(d, H), _linear(d, H)
        self.A_log, self.dt_bias = nn.Parameter(torch.zeros(H)), nn.Parameter(torch.zeros(H))
        self.q_conv1d, self.k_conv1d, self.v_conv1d = _Weight(H * dk, 1, K), _Weight(H * dk, 1, K), _Weight(H * dv, 1, K)
        self.o_norm = _Weight(dv)


class SwiGLU(nn.Module):
    def __init__(self, d, f):
        super().__init__()
        self.w_gate, self.w_up, self.w_down = _linear(d, f), _linear(d, f), _linear(f, d)


class Block(nn.Module):
    def __init__(self, c, i, has_attn_sources):
        super().__init__()
        self.is_attn = (i + 1) % c["full_attn_every"] == 0
        d = c["d_model"]
        self.norm1, self.norm2 = _Weight(d), _Weight(d)
        self.mixer = GatedAttention(c) if self.is_attn else GatedDeltaNet(c)
        self.mlp = SwiGLU(d, c["ffn_dim"])
        self.res1, self.res2 = _Gate(d), _Gate(d)
        self.delta_mlp_query, self.delta_mlp_weight = nn.Parameter(torch.zeros(d)), nn.Parameter(torch.ones(d))
        if has_attn_sources:
            self.delta_attn_query, self.delta_attn_weight = nn.Parameter(torch.zeros(d)), nn.Parameter(torch.ones(d))


class NgramMemory(nn.Module):
    def __init__(self, ng, d_model):
        super().__init__()
        T, D, n = len(ng["orders"]) * ng["heads_per_order"], ng["dim"], len(ng["inject_layers"])
        self.proj = nn.ModuleList([_linear(T * D, d_model) for _ in range(n)])
        self.key_norm = nn.ModuleList([_Weight(d_model) for _ in range(n)])
        self.val_norm = nn.ModuleList([_Weight(d_model) for _ in range(n)])
        self.register_buffer("sign_pool", torch.zeros(ng["sign_pool"], D), persistent=True)
        # table [T * rows, D] (BF16) and the canonical-id map are attached by from_pretrained (not in model.safetensors)



class _State:
    """Static decode buffers (one active sequence). Fixed addresses so the decode step can be a CUDA graph."""

    def __init__(self, m, max_len, device):
        self.max_len, self.device = max_len, device
        C = 2 * m.GH * m.GK + m.GH * m.GV
        self.conv = [None if a else torch.zeros(C, m.CONV - 1, device=device) for a in m.is_attn]
        self.S = [None if a else torch.zeros(1, m.GH, m.GK, m.GV, device=device) for a in m.is_attn]
        self.k = [torch.zeros(1, m.KV, max_len, m.HD, device=device, dtype=m.mm_dtype) if a else None for a in m.is_attn]
        self.v = [torch.zeros(1, m.KV, max_len, m.HD, device=device, dtype=m.mm_dtype) if a else None for a in m.is_attn]
        self.ctx = torch.zeros(3, dtype=torch.int64, device=device)    # last 3 canonical ids (trigram window)
        self.pos = torch.zeros(1, dtype=torch.int64, device=device)    # next position
        self.tok = torch.zeros(1, dtype=torch.int64, device=device)    # decode-step input
        self.graph = self.logits = None
        self.owner = None

    def recurrent(self):
        return [t for t in self.conv + self.S + [self.ctx, self.pos] if t is not None]


class HybridCache:
    """What generate() passes back as past_key_values. A new prefill takes over the model's static buffers."""
    is_compileable = False

    def __init__(self, state):
        self.state, self.seen = state, 0
        state.owner = self

    def get_seq_length(self, layer_idx=0):
        return self.seen


# ----------------------------------------------------------------------------- model
class HybridLMForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = HybridLMConfig
    base_model_prefix = "hybrid"
    _no_split_modules = ["Block"]
    _is_stateful = True
    _supports_sdpa = True
    # kept in FP32 even with dtype=torch.bfloat16 (training's residual stream / norms / gates were FP32)
    _keep_in_fp32_modules_strict = ["norm", "res1", "res2", "delta_", "A_log", "dt_bias", "embed", "sign_pool",
                                    "conv1d"]

    def __init__(self, config):
        super().__init__(config)
        c, mods = config.model_config, config.modifications
        assert mods.get("delta_block") and mods.get("delta_variant", "repo") == "repo"
        assert mods.get("depth_attn", "off") == "off"
        self.D, self.EPS = c["d_model"], c["norm_eps"]
        self.H, self.KV, self.HD = c["n_heads"], c["n_kv_heads"], c["head_dim"]
        self.GH, self.GK = c["gdn_heads"], c["gdn_head_dim"]
        self.GV, self.CONV = int(self.GK * c["gdn_expand_v"]), c["gdn_conv_size"]
        n = c["n_layers"]
        self.per_block = math.ceil(n / mods.get("delta_blocks", 4))
        self.is_attn = [(i + 1) % c["full_attn_every"] == 0 for i in range(n)]
        ng = c["ngram"]
        self.ng = ng
        self.rows = max(1024, ng["total_params"] // (len(ng["orders"]) * ng["heads_per_order"] * ng["dim"]))
        self.inject = {l: k for k, l in enumerate(ng["inject_layers"])}
        self.embed = nn.Embedding(c["vocab_size"], c["d_model"])
        # every layer after the first sees at least one Delta Block source (the embedding) at its mixer input
        self.blocks = nn.ModuleList([Block(c, i, i > 0) for i in range(n)])
        self.norm = _Weight(c["d_model"])
        self.lm_head = _linear(c["d_model"], c["vocab_size"])
        self.ngram = NgramMemory(ng, c["d_model"])
        self._fused_ptr, self._state = None, None
        self.post_init()

    def _init_weights(self, module):
        pass    # everything comes from the checkpoint

    def get_input_embeddings(self):
        return self.embed

    def get_output_embeddings(self):
        return self.lm_head

    @classmethod
    def _supports_default_dynamic_cache(cls):
        return False    # HybridCache is created by forward()

    # ------------------------------------------------------------------ loading
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
        from safetensors import safe_open
        from transformers.utils import cached_file
        hub = {k: kwargs[k] for k in ("revision", "token", "cache_dir", "local_files_only") if k in kwargs}
        device = model.embed.weight.device
        with safe_open(cached_file(pretrained_model_name_or_path, "ngram.safetensors", **hub), "pt", device="cpu") as f:
            table = f.get_tensor("table")
        expected = (len(model.ng["orders"]) * model.ng["heads_per_order"] * model.rows, model.ng["dim"])
        assert tuple(table.shape) == expected, (table.shape, expected)
        canon = build_canon_map(cached_file(pretrained_model_name_or_path, "tokenizer.json", **hub))
        sha = model.config.modifications.get("ngram_canonical_sha256")
        if model.ng.get("canonical_tokens"):
            assert sha is None or hashlib.sha256(canon.numpy().tobytes()).hexdigest() == sha, \
                "canonical-token map differs from the one used in training"
        else:
            canon = torch.arange(model.config.vocab_size)
        model.ngram.register_buffer("table", table.to(device), persistent=False)
        model.ngram.register_buffer("canon", canon.to(device), persistent=False)
        model.eval()
        return model

    @property
    def mm_dtype(self):
        return self.lm_head.weight.dtype

    def _ready(self):
        """Fuse projections into single GEMMs (the original parameters become views into the fused weight)."""
        if self._fused_ptr == self.blocks[0].mixer.q_proj.weight.data_ptr():
            return
        with torch.no_grad():
            for blk in self.blocks:
                m = blk.mixer
                names = (("q_proj", "k_proj", "v_proj", "gate_proj") if blk.is_attn else
                         ("q_proj", "k_proj", "v_proj", "a_proj", "b_proj", "g_proj"))
                blk._in_w = self._fuse([getattr(m, n) for n in names])
                blk._mlp_w = self._fuse([blk.mlp.w_gate, blk.mlp.w_up])
                if not blk.is_attn:
                    blk._conv_w = torch.cat([m.q_conv1d.weight[:, 0], m.k_conv1d.weight[:, 0],
                                             m.v_conv1d.weight[:, 0]]).float()
                    blk._neg_A = -m.A_log.float().exp()
                blk._g1, blk._g2 = 2 * torch.sigmoid(blk.res1.a.float()), 2 * torch.sigmoid(blk.res2.a.float())
                blk._dmq = (blk.delta_mlp_query * blk.delta_mlp_weight).float()
                if hasattr(blk, "delta_attn_query"):
                    blk._daq = (blk.delta_attn_query * blk.delta_attn_weight).float()
        self._fused_ptr = self.blocks[0].mixer.q_proj.weight.data_ptr()
        self._state = None

    @staticmethod
    def _fuse(linears):
        w = torch.cat([l.weight.data for l in linears])
        o = 0
        for l in linears:
            n = l.weight.shape[0]
            l.weight.data = w[o:o + n]
            o += n
        return w

    def _get_state(self):
        dev = self.embed.weight.device
        if self._state is None or self._state.device != dev:
            self._state = _State(self, self.config.kv_cache_len, dev)
        return self._state

    # ------------------------------------------------------------------ pieces
    def _ngram_memory(self, cids, pos):
        ng, rows, signs, valid = self.ng, [], [], []
        t = 0
        for order in ng["orders"]:
            for head in range(ng["heads_per_order"]):
                h = torch.zeros_like(cids)
                m = _i64(_MULTS[(head * 3 + order) % 8])
                for j in range(order):
                    h = h * m + (_shift(cids, j) + 1 + j * 7919)
                h = h ^ (h >> 29); h = h * _i64(0xBF58476D1CE4E5B9); h = h ^ (h >> 32)
                rows.append(torch.remainder(h, self.rows) + t * self.rows)
                if ng.get("sign_pool"):
                    s = torch.full_like(cids, _i64(0x632BE59BD9B4E019))
                    m = _i64(_SIGN_MULTS[(head * 3 + order) % 8])
                    for j in range(order):
                        s = s * m + (_shift(cids, j) + 1 + j * 104729)
                    s = s ^ (s >> 31); s = s * _i64(0x94D049BB133111EB); s = s ^ (s >> 29)
                    signs.append(s & (ng["sign_pool"] - 1))
                valid.append(pos >= order - 1)
                t += 1
        emb = self.ngram.table[torch.stack(rows, 1)].float()                        # [N, T, dim]
        mult = torch.stack(valid, 1).unsqueeze(-1).float()
        if signs:
            mult = self.ngram.sign_pool[torch.stack(signs, 1)] * mult
        return (emb * mult).flatten(1)

    def _rope(self, x, pos):
        d = x.shape[-1]
        inv = 1.0 / (self.config.model_config["rope_theta"] ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d))
        f = pos.float()[:, None] * inv[None]
        cos, sin = f.cos()[:, None], f.sin()[:, None]
        x1, x2 = x[..., : d // 2], x[..., d // 2:]
        return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1)

    def _gdn(self, blk, i, xn, st, decode):
        m, N, H, dk, dv = blk.mixer, xn.shape[0], self.GH, self.GK, self.GV
        qkv, a, b, gate = (xn @ blk._in_w.T).split((2 * H * dk + H * dv, H, H, H * dv), -1)
        w = blk._conv_w
        if decode:                      # one causal depthwise-conv step from the cached last K-1 inputs
            win = torch.cat([st.conv[i], qkv.T.float()], 1)
            y = (win * w).sum(-1, keepdim=True)
            st.conv[i].copy_(win[:, 1:])
        else:
            win = F.pad(qkv.T.float(), (self.CONV - 1, 0))
            y = F.conv1d(win[None], w[:, None], groups=w.shape[0])[0]
            if st is not None:
                st.conv[i].copy_(win[:, -(self.CONV - 1):])
        q, k, v = F.silu(y).T.to(xn.dtype).split((H * dk, H * dk, H * dv), -1)
        g = (blk._neg_A * F.softplus(a.float() + m.dt_bias.float()))[None]         # log decay
        beta = b.float().sigmoid()[None]
        kern = fused_recurrent_gated_delta_rule if decode else chunk_gated_delta_rule
        o, S = kern(q=q.reshape(1, N, H, dk), k=k.reshape(1, N, H, dk), v=v.reshape(1, N, H, dv), g=g, beta=beta,
                    scale=dk ** -0.5, initial_state=st.S[i] if decode else None,
                    output_final_state=st is not None, use_qk_l2norm_in_kernel=True)
        if st is not None:
            st.S[i].copy_(S)
        o = _rms(o[0], m.o_norm.weight, self.EPS) * F.silu(gate.float().view(N, H, dv))
        return o.reshape(N, -1).to(xn.dtype) @ m.o_proj.weight.T

    def _attention(self, blk, i, xn, pos, st, decode):
        m, N, H, KV, HD = blk.mixer, xn.shape[0], self.H, self.KV, self.HD
        q, k, v, gate = (xn @ blk._in_w.T).split((H * HD, KV * HD, KV * HD, H * HD), -1)
        q = self._rope(_rms(q.view(N, H, HD), m.q_norm.weight, self.EPS), pos).to(xn.dtype)
        k = self._rope(_rms(k.view(N, KV, HD), m.k_norm.weight, self.EPS), pos).to(xn.dtype)
        v = v.view(N, KV, HD)
        if decode:
            st.k[i].index_copy_(2, pos, k.transpose(0, 1)[None])
            st.v[i].index_copy_(2, pos, v.transpose(0, 1)[None])
            mask = (torch.arange(st.max_len, device=xn.device) <= pos)[None, None, None]
            o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], st.k[i], st.v[i], attn_mask=mask,
                                               enable_gqa=True)
        else:
            if st is not None:
                st.k[i][:, :, :N] = k.transpose(0, 1)[None]
                st.v[i][:, :, :N] = v.transpose(0, 1)[None]
            o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], k.transpose(0, 1)[None],
                                               v.transpose(0, 1)[None], is_causal=True, enable_gqa=True)
        o = o[0].transpose(0, 1) * torch.sigmoid(gate.view(N, H, HD))
        return o.reshape(N, -1) @ m.o_proj.weight.T

    def _run(self, ids, st, decode, keep):
        """ids [N] (prefill from position 0) or [1] (decode at st.pos). Returns logits [keep or N, V] (FP32)."""
        N, dt, eps = ids.shape[0], self.mm_dtype, self.EPS
        cids = self.ngram.canon[ids]
        if decode:
            pos = st.pos
            st.ctx.copy_(torch.cat([st.ctx[1:], cids]))
            mem = self._ngram_memory(st.ctx, torch.cat([pos - 2, pos - 1, pos]))[-1:]
        else:
            pos = torch.arange(N, device=ids.device)
            mem = self._ngram_memory(cids, pos)
            if st is not None:
                st.ctx.copy_(F.pad(cids, (3, 0))[-3:])
                st.pos.fill_(N)
        x = self.embed.weight[ids].float()                 # residual stream in FP32
        sources = []
        for i, blk in enumerate(self.blocks):
            if i in self.inject:                           # n-gram memory, context-gated add
                k = self.inject[i]
                e = (mem.to(dt) @ self.ngram.proj[k].weight.T).float()
                gate = torch.sigmoid((_rms(x, self.ngram.key_norm[k].weight, eps) *
                                      _rms(e, self.ngram.val_norm[k].weight, eps)).sum(-1, keepdim=True) / math.sqrt(self.D))
                x = x + gate * e
            attn_sources = list(sources)
            if i % self.per_block == 0:                    # Delta Block source, "repo": s_k = x_k - s_{k-1}
                sources.append(x if not sources else x - sources[-1])
            h = x + _select(attn_sources, blk._daq, eps) if attn_sources else x
            xn = _rms(h, blk.norm1.weight, eps).to(dt)
            y = self._attention(blk, i, xn, pos, st, decode) if blk.is_attn else self._gdn(blk, i, xn, st, decode)
            x = x + blk._g1 * y.float()
            h = x + _select(sources, blk._dmq, eps)
            gt, up = (_rms(h, blk.norm2.weight, eps).to(dt) @ blk._mlp_w.T).chunk(2, -1)
            x = x + blk._g2 * ((F.silu(gt) * up) @ blk.mlp.w_down.weight.T).float()
        if keep:
            x = x[-keep:]
        logits = (_rms(x, self.norm.weight, eps).to(dt) @ self.lm_head.weight.T).float()
        if decode:
            st.pos += 1
        return logits

    def _decode(self, st, token):
        st.tok.copy_(token)
        if not (self.config.use_cuda_graph and st.device.type == "cuda"):
            return self._run(st.tok, st, True, 0)
        if st.graph is None:
            snap = [t.clone() for t in st.recurrent()]
            s = torch.cuda.Stream()
            s.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(s):                     # warm up (Triton autotune) outside the graph
                for _ in range(2):
                    self._run(st.tok, st, True, 0)
            torch.cuda.current_stream().wait_stream(s)
            graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(graph):
                st.logits = self._run(st.tok, st, True, 0)
            for dst, src in zip(st.recurrent(), snap):
                dst.copy_(src)
            st.graph = graph
        st.graph.replay()
        return st.logits.clone()

    # ------------------------------------------------------------------ HF API
    def forward(self, input_ids=None, attention_mask=None, past_key_values=None, use_cache=None, labels=None,
                logits_to_keep=0, return_dict=True, **kwargs):
        if input_ids is None or input_ids.shape[0] != 1:
            raise ValueError("this model implementation supports input_ids of batch size 1")
        if attention_mask is not None and not bool(attention_mask.all()):
            raise ValueError("padding is not supported (batch size 1)")
        self._ready()
        use_cache = self.config.use_cache if use_cache is None else use_cache
        ids = input_ids[0]
        keep = logits_to_keep if isinstance(logits_to_keep, int) else 0
        cache = past_key_values if isinstance(past_key_values, HybridCache) else None
        with torch.no_grad():
            if cache is not None and cache.seen > 0:       # continue a cached sequence token by token
                assert cache.state.owner is cache, "this cache was superseded by a newer prefill"
                if cache.seen + ids.numel() > cache.state.max_len:
                    raise ValueError(f"sequence exceeds kv_cache_len={cache.state.max_len}")
                logits = torch.cat([self._decode(cache.state, t) for t in ids.split(1)])
                cache.seen += ids.numel()
            else:
                cache = None
                if use_cache:
                    if ids.numel() > self.config.kv_cache_len:
                        raise ValueError(f"prompt exceeds kv_cache_len={self.config.kv_cache_len}")
                    cache = HybridCache(self._get_state())
                    cache.seen = ids.numel()
                logits = self._run(ids, cache.state if cache else None, False, 0 if labels is not None else keep)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[:-1], labels[0, 1:], ignore_index=-100)
        return CausalLMOutputWithPast(loss=loss, logits=logits[None], past_key_values=cache)

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, use_cache=True,
                                      **kwargs):
        if isinstance(past_key_values, HybridCache) and past_key_values.seen > 0:
            input_ids = input_ids[:, past_key_values.seen:]
        return dict(input_ids=input_ids, past_key_values=past_key_values, use_cache=use_cache, logits_to_keep=1)


HybridLMConfig.register_for_auto_class()
HybridLMForCausalLM.register_for_auto_class("AutoModelForCausalLM")