Fix .generate() support: add GenerationMixin + generation_config, document KV-cache limitation

modeling_nanochat.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers import PreTrainedModel
+from transformers import PreTrainedModel, GenerationMixin, GenerationConfig
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from .configuration_nanochat import NanochatConfig
@@ -40,16 +40,13 @@ def _apply_rotary(x, cos, sin):
 
 
 def _sdpa_attention(q, k, v, window_size, enable_gqa):
-    # q/k/v are (B, H, T, D)
     t_q = q.size(2)
     t_k = k.size(2)
     left_window = window_size[0]
 
-    # Full causal attention when the window covers full context.
     if (left_window < 0 or left_window >= t_q) and t_q == t_k:
         return F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=enable_gqa)
 
-    # Single-token decode path.
     if t_q == 1:
         if left_window >= 0 and left_window < t_k:
             start = max(0, t_k - (left_window + 1))
@@ -57,7 +54,6 @@ def _sdpa_attention(q, k, v, window_size, enable_gqa):
             v = v[:, :, start:, :]
         return F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa)
 
-    # Build explicit causal (+ optional sliding-window) mask.
     device = q.device
     row_idx = (t_k - t_q) + torch.arange(t_q, device=device).unsqueeze(1)
     col_idx = torch.arange(t_k, device=device).unsqueeze(0)
@@ -71,7 +67,6 @@ def _sdpa_attention(q, k, v, window_size, enable_gqa):
 def _flash_attn_func(q, k, v, causal=False, window_size=(-1, -1)):
     if not causal:
         raise NotImplementedError("Nanochat HF export currently supports only causal attention")
-    # SDPA fallback mirroring nanochat.flash_attention semantics.
     q = q.transpose(1, 2)
     k = k.transpose(1, 2)
     v = v.transpose(1, 2)
@@ -240,7 +235,7 @@ class NanochatBackbone(nn.Module):
         return logits
 
 
-class NanochatForCausalLM(PreTrainedModel):
+class NanochatForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = NanochatConfig
     base_model_prefix = "model"
     main_input_name = "input_ids"
@@ -249,6 +244,8 @@ class NanochatForCausalLM(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.model = NanochatBackbone(config)
+        if getattr(self, "generation_config", None) is None:
+            self.generation_config = GenerationConfig.from_model_config(config)
 
     @property
     def all_tied_weights_keys(self):
@@ -299,4 +296,4 @@ class NanochatForCausalLM(PreTrainedModel):
         )
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-        return {"input_ids": input_ids}
+        return {"input_ids": input_ids}