Update modeling_spin.py

modeling_spin.py

@@ -6,45 +6,6 @@ from transformers import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from .configuration_spin import SpinConfig
 
-try:
-    from transformers.cache_utils import DynamicCache
-    HAVE_DYNAMIC_CACHE = True
-except ImportError:
-    HAVE_DYNAMIC_CACHE = False
-
-
-def extract_kv_cache(past_key_values):
-    """Safely extracts KV cache regardless of whether HF passes DynamicCache, list, or tuple."""
-    if past_key_values is None:
-        return None, 0
-
-    # 1. Hugging Face DynamicCache / Cache object
-    if hasattr(past_key_values, "to_legacy_cache"):
-        legacy = past_key_values.to_legacy_cache()
-        seq_len = past_key_values.get_seq_length() if hasattr(past_key_values, "get_seq_length") else 0
-        if seq_len == 0 or len(legacy) == 0:
-            return None, 0
-        return list(legacy), seq_len
-
-    if hasattr(past_key_values, "key_cache") and hasattr(past_key_values, "value_cache"):
-        if len(past_key_values.key_cache) == 0:
-            return None, 0
-        kv_list = [
-            (past_key_values.key_cache[i], past_key_values.value_cache[i])
-            for i in range(len(past_key_values.key_cache))
-        ]
-        seq_len = kv_list[0][0].shape[2] if len(kv_list) > 0 else 0
-        return kv_list, seq_len
-
-    # 2. Legacy tuple / list format
-    if isinstance(past_key_values, (list, tuple)) and len(past_key_values) > 0:
-        if past_key_values[0] is None:
-            return None, 0
-        seq_len = past_key_values[0][0].shape[2]
-        return list(past_key_values), seq_len
-
-    return None, 0
-
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-5):
@@ -93,8 +54,9 @@ class SwiGLU(nn.Module):
 
 
 class CausalSelfAttention(nn.Module):
-    def __init__(self, config: SpinConfig):
+    def __init__(self, config: SpinConfig, layer_idx: int = 0):
         super().__init__()
+        self.layer_idx = layer_idx
         self.n_heads = config.n_heads
         self.head_dim = config.d_model // config.n_heads
 
@@ -106,7 +68,7 @@ class CausalSelfAttention(nn.Module):
         mask = torch.full((config.max_seq_len, config.max_seq_len), float("-inf"))
         self.register_buffer("causal_mask", torch.triu(mask, diagonal=1), persistent=False)
 
-    def forward(self, x, freqs_cos, freqs_sin, kv_cache=None):
+    def forward(self, x, freqs_cos, freqs_sin, past_key_value=None):
         B, T, C = x.shape
         q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
         k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim)
@@ -115,11 +77,20 @@ class CausalSelfAttention(nn.Module):
         q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
         q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
 
-        if kv_cache is not None:
-            prev_k, prev_v = kv_cache
-            k = torch.cat([prev_k, k], dim=2)
-            v = torch.cat([prev_v, v], dim=2)
-        new_kv_cache = (k, v)
+        # Standard Cache update (handles both DynamicCache and classic tuple)
+        if past_key_value is not None:
+            if hasattr(past_key_value, "update"):
+                k, v = past_key_value.update(k, v, self.layer_idx)
+                new_kv_cache = past_key_value
+            elif isinstance(past_key_value, tuple):
+                prev_k, prev_v = past_key_value
+                k = torch.cat([prev_k, k], dim=2)
+                v = torch.cat([prev_v, v], dim=2)
+                new_kv_cache = (k, v)
+            else:
+                new_kv_cache = (k, v)
+        else:
+            new_kv_cache = (k, v)
 
         scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
         if T > 1:
@@ -131,15 +102,15 @@ class CausalSelfAttention(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, config: SpinConfig):
+    def __init__(self, config: SpinConfig, layer_idx: int = 0):
         super().__init__()
         self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
-        self.attn = CausalSelfAttention(config)
+        self.attn = CausalSelfAttention(config, layer_idx=layer_idx)
         self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
         self.ffn = SwiGLU(config.d_model, config.d_ff)
 
-    def forward(self, x, freqs_cos, freqs_sin, kv_cache=None):
-        attn_out, next_kv = self.attn(self.attn_norm(x), freqs_cos, freqs_sin, kv_cache=kv_cache)
+    def forward(self, x, freqs_cos, freqs_sin, past_key_value=None):
+        attn_out, next_kv = self.attn(self.attn_norm(x), freqs_cos, freqs_sin, past_key_value=past_key_value)
         x = x + attn_out
         x = x + self.ffn(self.ffn_norm(x))
         return x, next_kv
@@ -148,12 +119,13 @@ class TransformerBlock(nn.Module):
 class SpinForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = SpinConfig
     _tied_weights_keys = {"lm_head.weight": "tok_embeddings.weight"}
+    _supports_cache_class = True
 
     def __init__(self, config: SpinConfig):
         super().__init__(config)
         self.config = config
         self.tok_embeddings = nn.Embedding(config.vocab_size, config.d_model)
-        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.layers = nn.ModuleList([TransformerBlock(config, layer_idx=i) for i in range(config.n_layers)])
         self.norm = RMSNorm(config.d_model, eps=config.norm_eps)
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
 
@@ -192,16 +164,29 @@ class SpinForCausalLM(PreTrainedModel, GenerationMixin):
         B, T = input_ids.shape
         x = self.tok_embeddings(input_ids)
 
-        kv_caches_in, start_pos = extract_kv_cache(past_key_values)
+        # Calculate start position for RoPE
+        start_pos = 0
+        if past_key_values is not None:
+            if hasattr(past_key_values, "get_seq_length"):
+                start_pos = past_key_values.get_seq_length()
+            elif isinstance(past_key_values, (tuple, list)) and len(past_key_values) > 0 and past_key_values[0] is not None:
+                start_pos = past_key_values[0][0].shape[2]
 
         freqs_cos = self.freqs_cos[start_pos : start_pos + T]
         freqs_sin = self.freqs_sin[start_pos : start_pos + T]
 
-        new_kv_caches = []
+        legacy_kv_caches = []
         for i, layer in enumerate(self.layers):
-            cache_i = kv_caches_in[i] if kv_caches_in is not None else None
-            x, new_cache = layer(x, freqs_cos, freqs_sin, kv_cache=cache_i)
-            new_kv_caches.append(new_cache)
+            if hasattr(past_key_values, "update"):
+                layer_cache = past_key_values
+            elif isinstance(past_key_values, (tuple, list)) and len(past_key_values) > i:
+                layer_cache = past_key_values[i]
+            else:
+                layer_cache = None
+
+            x, new_cache = layer(x, freqs_cos, freqs_sin, past_key_value=layer_cache)
+            if not hasattr(past_key_values, "update"):
+                legacy_kv_caches.append(new_cache)
 
         x = self.norm(x)
         logits = self.lm_head(x)
@@ -210,32 +195,34 @@ class SpinForCausalLM(PreTrainedModel, GenerationMixin):
         if labels is not None:
             loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100)
 
-        # Package past_key_values back into HF expected format
         if use_cache:
-            if HAVE_DYNAMIC_CACHE and isinstance(past_key_values, DynamicCache):
-                past_key_values_out = DynamicCache.from_legacy_cache(new_kv_caches)
-            else:
-                past_key_values_out = tuple(new_kv_caches)
+            output_cache = past_key_values if hasattr(past_key_values, "update") else tuple(legacy_kv_caches)
         else:
-            past_key_values_out = None
+            output_cache = None
 
         if not return_dict:
-            return (logits, loss, past_key_values_out)
+            return (logits, loss, output_cache)
 
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
-            past_key_values=past_key_values_out,
+            past_key_values=output_cache,
         )
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
-        _, seq_len = extract_kv_cache(past_key_values)
-        if seq_len > 0:
+        past_length = 0
+        if past_key_values is not None:
+            if hasattr(past_key_values, "get_seq_length"):
+                past_length = past_key_values.get_seq_length()
+            elif isinstance(past_key_values, (tuple, list)) and len(past_key_values) > 0 and past_key_values[0] is not None:
+                past_length = past_key_values[0][0].shape[2]
+
+        if past_length > 0:
             input_ids = input_ids[:, -1:]
+
         return {
             "input_ids": input_ids,
             "past_key_values": past_key_values,
             "attention_mask": attention_mask,
             "use_cache": True,
         }
-