Upload NSAForCausalLM

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,36 @@
+{
+  "architectures": [
+    "NSAForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_nsa.NSAConfig",
+    "AutoModelForCausalLM": "modeling_nsa.NSAForCausalLM"
+  },
+  "d_k": 64,
+  "d_v": 64,
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "hidden_size": 768,
+  "max_position_embeddings": 2048,
+  "model_type": "nsa",
+  "n_kv_groups": 2,
+  "nsa": {
+    "block": 32,
+    "branches": [
+      "cmp",
+      "sel",
+      "win"
+    ],
+    "gqa_groups": 2,
+    "sel_block": 64,
+    "sel_top_n": 16,
+    "stride": 16,
+    "window": 512
+  },
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "rope_theta": 10000,
+  "transformers_version": "4.56.0",
+  "vocab_size": 256
+}

configuration_nsa.py

@@ -0,0 +1,40 @@
+# Remote code: configuration and modeling for NSA
+from transformers import PretrainedConfig
+
+
+class NSAConfig(PretrainedConfig):
+    model_type = "nsa"
+
+    def __init__(
+        self,
+        vocab_size=50257,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        n_kv_groups=1,
+        d_k=64,
+        d_v=64,
+        max_position_embeddings=2048,
+        rope_theta=10000,
+        nsa=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.n_kv_groups = n_kv_groups
+        self.d_k = d_k
+        self.d_v = d_v
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.nsa = nsa or {
+            "branches": ["cmp", "sel", "win"],
+            "window": 512,
+            "gqa_groups": n_kv_groups,
+            "block": 32,
+            "stride": 16,
+            "sel_block": 64,
+            "sel_top_n": 16,
+        }

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [
+    0
+  ],
+  "pad_token_id": 0,
+  "transformers_version": "4.56.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a286e8c0a9e11a7bd3355d520a92a32b849212de186aa5c83c19ca23bcda7d8
+size 313204760

modeling_nsa.py

@@ -0,0 +1,329 @@
+# Remote code: configuration and modeling for NSA
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutput
+
+from .configuration_nsa import NSAConfig
+_HAS_NSA = False  # Do not attempt nested vendor import in HF dynamic loader
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
+        return (x * rms) * self.weight
+
+
+class MLP(nn.Module):
+    def __init__(self, dim: int, hidden_mult: int = 4) -> None:
+        super().__init__()
+        h = hidden_mult * dim
+        self.fc1 = nn.Linear(dim, h, bias=False)
+        self.fc2 = nn.Linear(h, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(torch.nn.functional.silu(self.fc1(x)))
+
+
+def _rope(q: torch.Tensor) -> torch.Tensor:
+    B, S, D = q.shape[0], q.shape[2], q.shape[-1]
+    if D % 2 != 0:
+        return q
+    device = q.device
+    half = D // 2
+    pos = torch.arange(S, device=device).float().unsqueeze(-1)
+    inv_freq = 1.0 / (10000 ** (torch.arange(0, half, device=device).float() / half))
+    angles = pos * inv_freq
+    cos = angles.cos().view(1, 1, S, half)
+    sin = angles.sin().view(1, 1, S, half)
+    q1, q2 = q[..., :half], q[..., half:]
+    return torch.cat([q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1)
+
+
+def _avg_pool_time(x: torch.Tensor, kernel: int, stride: int) -> torch.Tensor:
+    if x.shape[2] < kernel:
+        return x[..., :0, :]
+    xt = x.permute(0, 3, 1, 2).contiguous()
+    y = torch.nn.functional.avg_pool2d(xt, kernel_size=(1, kernel), stride=(1, stride))
+    return y.permute(0, 2, 3, 1).contiguous()
+
+
+def _window_mask(q: torch.Tensor, S: int, w: int) -> torch.Tensor:
+    B, h = q.shape[0], q.shape[1]
+    device = q.device
+    row = torch.arange(S, device=device).view(S, 1)
+    col = torch.arange(S, device=device).view(1, S)
+    allowed = (col <= row) & (col >= (row - (w - 1)))
+    M = torch.full((S, S), float('-inf'), device=device, dtype=q.dtype)
+    M.masked_fill_(allowed, 0.0)
+    return M.view(1, 1, S, S).expand(B, h, S, S)
+
+
+def _selection_blocks(scores: torch.Tensor, l_sel: int, n_sel: int) -> torch.Tensor:
+    B, h, S = scores.shape
+    n_blocks = max(1, (S + l_sel - 1) // l_sel)
+    # Pad to multiple of l_sel
+    pad = n_blocks * l_sel - S
+    if pad > 0:
+        scores = torch.nn.functional.pad(scores, (0, pad), value=-1e9)
+    blk_scores = scores.view(B, h, n_blocks, l_sel).max(dim=-1).values
+    k = min(n_sel, n_blocks)
+    return torch.topk(blk_scores, k=k, dim=-1).indices
+
+
+class EmbeddedNSAAttention(nn.Module):
+    def __init__(self, dim: int, n_heads: int, n_kv_groups: int, d_k: int, d_v: int,
+                 l: int, d: int, l_sel: int, n_sel: int, w: int) -> None:
+        super().__init__()
+        self.n_heads = n_heads
+        self.n_kv_groups = n_kv_groups
+        self.d_k = d_k
+        self.d_v = d_v
+        self.l = l
+        self.stride = d
+        self.l_sel = l_sel
+        self.n_sel = n_sel
+        self.w = w
+        self.W_Q = nn.Linear(dim, n_heads * d_k, bias=False)
+        self.W_K_cmp = nn.Linear(dim, n_kv_groups * d_k, bias=False)
+        self.W_V_cmp = nn.Linear(dim, n_kv_groups * d_v, bias=False)
+        self.W_K_sel = nn.Linear(dim, n_kv_groups * d_k, bias=False)
+        self.W_V_sel = nn.Linear(dim, n_kv_groups * d_v, bias=False)
+        self.W_K_win = nn.Linear(dim, n_kv_groups * d_k, bias=False)
+        self.W_V_win = nn.Linear(dim, n_kv_groups * d_v, bias=False)
+        # Gate MLP operates on per-group pooled Q with width d_k (matches training)
+        gate_hidden = max(1, d_k // 2)
+        self.gate_fc1 = nn.Linear(d_k, gate_hidden, bias=True)
+        self.gate_fc2 = nn.Linear(gate_hidden, 3, bias=True)
+        nn.init.xavier_uniform_(self.gate_fc2.weight, gain=0.1)
+        nn.init.zeros_(self.gate_fc2.bias)
+        self.out = nn.Linear(n_heads * d_v, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, S, D = x.shape
+        h, dk, dv = self.n_heads, self.d_k, self.d_v
+        Q = self.W_Q(x).view(B, S, h, dk).transpose(1, 2)  # [B,h,S,dk]
+        g = max(1, self.n_kv_groups)
+        r = max(1, h // g)
+        # Project per-group K/V then broadcast to heads
+        Kc_g = self.W_K_cmp(x).view(B, S, g, dk).permute(0, 2, 1, 3)  # [B,g,S,dk]
+        Vc_g = self.W_V_cmp(x).view(B, S, g, dv).permute(0, 2, 1, 3)
+        Ks_g = self.W_K_sel(x).view(B, S, g, dk).permute(0, 2, 1, 3)
+        Vs_g = self.W_V_sel(x).view(B, S, g, dv).permute(0, 2, 1, 3)
+        Kw_g = self.W_K_win(x).view(B, S, g, dk).permute(0, 2, 1, 3)
+        Vw_g = self.W_V_win(x).view(B, S, g, dv).permute(0, 2, 1, 3)
+        # Broadcast groups to heads
+        def _bcast_to_heads(T):
+            return T.unsqueeze(1).expand(B, r, g, S, T.shape[-1]).reshape(B, h, S, T.shape[-1])
+        Kc = _bcast_to_heads(Kc_g)
+        Vc = _bcast_to_heads(Vc_g)
+        Ks = _bcast_to_heads(Ks_g)
+        Vs = _bcast_to_heads(Vs_g)
+        Kw = _bcast_to_heads(Kw_g)
+        Vw = _bcast_to_heads(Vw_g)
+
+        # RoPE
+        Qr = _rope(Q.transpose(1, 2)).transpose(1, 2)
+        Kc_r = _rope(Kc.transpose(1, 2)).transpose(1, 2)
+        Ks_r = _rope(Ks.transpose(1, 2)).transpose(1, 2)
+        Kw_r = _rope(Kw.transpose(1, 2)).transpose(1, 2)
+
+        # Compressed: average-pool along time
+        Kc_p = _avg_pool_time(Kc_r, kernel=max(1, self.stride), stride=max(1, self.stride))
+        Vc_p = _avg_pool_time(Vc, kernel=max(1, self.stride), stride=max(1, self.stride))
+        O_cmp = torch.nn.functional.scaled_dot_product_attention(Qr, Kc_p, Vc_p, is_causal=True)
+
+        # Selection: naive top-n blocks (global), enforce causal via triangular mask
+        scores = (Qr * Ks_r).mean(dim=-1)  # [B,h,S]
+        blk_idx = _selection_blocks(scores, self.l_sel, self.n_sel)  # [B,h,n]
+        n_blocks = max(1, (S + self.l_sel - 1) // self.l_sel)
+        keep = torch.zeros((B, h, n_blocks), device=x.device, dtype=torch.bool)
+        keep.scatter_(2, blk_idx, True)
+        keep = keep.unsqueeze(-1).expand(B, h, n_blocks, self.l_sel).reshape(B, h, -1)[:, :, :S]
+        logits = torch.matmul(Qr / math.sqrt(dk), Ks_r.transpose(-2, -1))  # [B,h,S,S]
+        tri = torch.triu(torch.ones((S, S), device=x.device, dtype=torch.bool), diagonal=1)
+        logits = logits.masked_fill(tri, float('-inf'))
+        sel_mask = torch.where(keep.unsqueeze(2).expand(B, h, S, S), torch.zeros((), device=x.device, dtype=Qr.dtype), torch.full((), float('-inf'), device=x.device, dtype=Qr.dtype))
+        P = torch.nn.functional.softmax(logits + sel_mask, dim=-1)
+        O_sel = torch.matmul(P, Vs)
+
+        # Sliding window
+        M = _window_mask(Qr, S, max(1, self.w))
+        logits_w = torch.matmul(Qr / math.sqrt(dk), Kw_r.transpose(-2, -1)) + M
+        P_w = torch.nn.functional.softmax(logits_w, dim=-1)
+        O_win = torch.matmul(P_w, Vw)
+
+        # Gate & mix: compute per-token, per-group gate from pooled Q
+        # Pool Q across heads within each kv-group
+        # Qr: [B,h,S,dk] -> reshape to [B,G,h_per_group,S,dk] then mean over h_per_group
+        G = max(1, self.n_kv_groups)
+        h_per_group = max(1, h // G)
+        Qg = Qr.view(B, G, h_per_group, S, dk).mean(dim=2)  # [B,G,S,dk]
+        Qg = Qg.permute(0, 2, 1, 3)  # [B,S,G,dk]
+        g1 = torch.nn.functional.silu(self.gate_fc1(Qg))
+        gate = torch.nn.functional.softmax(self.gate_fc2(g1), dim=-1)  # [B,S,G,3]
+        gc = gate[..., 0:1].unsqueeze(-1)  # [B,S,G,1,1]
+        gs = gate[..., 1:2].unsqueeze(-1)
+        gw = gate[..., 2:3].unsqueeze(-1)
+        # Broadcast group gates to heads within the group
+        # Reshape branch outputs to [B,S,G,h_per_group,dv]
+        Oc = O_cmp.permute(0,2,1,3).view(B, S, G, h_per_group, dv)
+        Os = O_sel.permute(0,2,1,3).view(B, S, G, h_per_group, dv)
+        Ow = O_win.permute(0,2,1,3).view(B, S, G, h_per_group, dv)
+        O = gc * Oc + gs * Os + gw * Ow
+        O = O.reshape(B, S, h, dv).permute(0, 2, 1, 3)
+        O = O.transpose(1, 2).reshape(B, S, h * dv)
+        return self.out(O)
+
+class SimpleAttention(nn.Module):
+    def __init__(self, dim: int, n_heads: int, d_k: int, d_v: int) -> None:
+        super().__init__()
+        self.n_heads = n_heads
+        self.d_k = d_k
+        self.d_v = d_v
+        self.q_proj = nn.Linear(dim, n_heads * d_k, bias=False)
+        self.k_proj = nn.Linear(dim, n_heads * d_k, bias=False)
+        self.v_proj = nn.Linear(dim, n_heads * d_v, bias=False)
+        self.out = nn.Linear(n_heads * d_v, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, S, D = x.shape
+        h, dk, dv = self.n_heads, self.d_k, self.d_v
+        q = self.q_proj(x).view(B, S, h, dk).transpose(1, 2)  # [B,h,S,dk]
+        k = self.k_proj(x).view(B, S, h, dk).transpose(1, 2)  # [B,h,S,dk]
+        v = self.v_proj(x).view(B, S, h, dv).transpose(1, 2)  # [B,h,S,dv]
+        attn = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
+        attn = attn.transpose(1, 2).contiguous().view(B, S, h * dv)
+        return self.out(attn)
+
+
+class SimpleBlock(nn.Module):
+    def __init__(self, dim: int, n_heads: int, d_k: int, d_v: int) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = SimpleAttention(dim, n_heads, d_k, d_v)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = MLP(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class NSABlockRemote(nn.Module):
+    """Transformer block with embedded NSA attention, pre/post RMSNorm, and MLP."""
+    def __init__(self, dim: int, n_heads: int, n_kv_groups: int, d_k: int, d_v: int,
+                 l: int, d: int, l_sel: int, n_sel: int, w: int) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = EmbeddedNSAAttention(dim, n_heads, n_kv_groups, d_k, d_v, l, d, l_sel, n_sel, w)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = MLP(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+class NSATinyLM(nn.Module):
+    def __init__(self, config: NSAConfig):
+        super().__init__()
+        self.config = config
+        self.vocab_size = int(config.vocab_size)
+        self.hidden_size = int(config.hidden_size)
+        self.num_hidden_layers = int(config.num_hidden_layers)
+        self.num_attention_heads = int(config.num_attention_heads)
+        self.n_kv_groups = int(getattr(config, "n_kv_groups", 1))
+        self.d_k = int(getattr(config, "d_k", self.hidden_size // self.num_attention_heads))
+        self.d_v = int(getattr(config, "d_v", self.hidden_size // self.num_attention_heads))
+        nsa = config.nsa or {}
+        self.l = int(nsa.get("block", 32))
+        self.d = int(nsa.get("stride", 16))
+        self.l_sel = int(nsa.get("sel_block", 64))
+        self.n_sel = int(nsa.get("sel_top_n", 16))
+        self.w = int(nsa.get("window", 512))
+
+        self.embed = nn.Embedding(self.vocab_size, self.hidden_size)
+        import os as _os
+        # Allow forcing simple fallback via env for integration tests
+        _force_simple = _os.getenv('NSA_REMOTE_FORCE_SIMPLE', '0').lower() in ('1','true','yes')
+        if not _force_simple:
+            # Fallback to embedded minimal NSA if vendor import failed
+            self.blocks = nn.ModuleList([
+                NSABlockRemote(
+                    self.hidden_size,
+                    self.num_attention_heads,
+                    self.n_kv_groups,
+                    self.d_k,
+                    self.d_v,
+                    self.l,
+                    self.d,
+                    self.l_sel,
+                    self.n_sel,
+                    self.w,
+                ) for _ in range(self.num_hidden_layers)
+            ])
+        else:
+            self.blocks = nn.ModuleList([
+                SimpleBlock(self.hidden_size, self.num_attention_heads, self.d_k, self.d_v)
+                for _ in range(self.num_hidden_layers)
+            ])
+        self.norm = nn.LayerNorm(self.hidden_size)
+        self.lm_head = nn.Linear(self.hidden_size, self.vocab_size, bias=False)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        x = self.embed(input_ids)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        return logits
+
+
+class NSAForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = NSAConfig
+    _no_split_modules = ["EmbeddedNSAAttention", "SimpleBlock"]
+
+    def __init__(self, config: NSAConfig):
+        super().__init__(config)
+        self.model = NSATinyLM(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed
+
+    def set_input_embeddings(self, new_emb):
+        self.model.embed = new_emb
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        if input_ids is None:
+            raise ValueError("input_ids is required")
+        logits = self.model(input_ids)
+        loss = None
+        if labels is not None:
+            # Shift for causal LM loss
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss_fct = torch.nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        # No past_key_values cache: rerun full sequence. Works everywhere, slower at decode.
+        return {"input_ids": input_ids, "attention_mask": kwargs.get("attention_mask", None)}