Update modeling_nandi.py

modeling_nandi.py

@@ -23,20 +23,20 @@ from collections.abc import Callable
 import torch
 import torch.nn as nn
 
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache, DynamicLayer
-from transformers.generation import GenerationMixin
-from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask
-from transformers.modeling_layers import GradientCheckpointingLayer
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
-from transformers.processing_utils import Unpack
-from transformers.utils import TransformersKwargs, auto_docstring
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple, merge_with_config_defaults
-from transformers.utils.output_capturing import capture_outputs
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, DynamicLayer
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ...utils.deprecation import deprecate_kwarg
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
 from .configuration_nandi import NandiConfig
 
 
@@ -189,7 +189,6 @@ class NandiAttention(nn.Module):
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None,
         past_key_values: Cache | None = None,
-        cache_position: torch.LongTensor | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         input_shape = hidden_states.shape[:-1]
@@ -203,8 +202,7 @@ class NandiAttention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
-            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
 
         attention_interface: Callable = eager_attention_forward
         if self.config._attn_implementation != "eager":
@@ -255,7 +253,6 @@ class NandiDecoderLayer(GradientCheckpointingLayer):
         position_ids: torch.LongTensor | None = None,
         past_key_values: Cache | None = None,
         use_cache: bool | None = False,
-        cache_position: torch.LongTensor | None = None,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> torch.Tensor:
@@ -268,7 +265,6 @@ class NandiDecoderLayer(GradientCheckpointingLayer):
             position_ids=position_ids,
             past_key_values=past_key_values,
             use_cache=use_cache,
-            cache_position=cache_position,
             position_embeddings=position_embeddings,
             **kwargs,
         )
@@ -354,7 +350,6 @@ class NandiModel(NandiPreTrainedModel):
         position_ids: torch.LongTensor | None = None,
         past_key_values: Cache | None = None,
         inputs_embeds: torch.FloatTensor | None = None,
-        cache_position: torch.LongTensor | None = None,
         use_cache: bool | None = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> BaseModelOutputWithPast:
@@ -370,20 +365,19 @@ class NandiModel(NandiPreTrainedModel):
         repeats = self.config.layer_sharing_repeats if self.config.layer_sharing else 1
 
         if use_cache and past_key_values is None:
+            # Use lazy DynamicCache (no config) so it grows to accommodate
+            # num_hidden_layers * repeats virtual slots for layer-sharing.
             past_key_values = DynamicCache()
 
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
-
         if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
 
         causal_mask = create_causal_mask(
             config=self.config,
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
-            cache_position=cache_position,
             past_key_values=past_key_values,
             position_ids=position_ids,
         )
@@ -393,6 +387,8 @@ class NandiModel(NandiPreTrainedModel):
 
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
             for repeat_idx in range(repeats):
+                # Each repeat gets its own virtual cache slots offset by num_hidden_layers,
+                # so repeat 0 uses slots 0..N-1 and repeat 1 uses slots N..2N-1, etc.
                 repeat_cache = (
                     _VirtualLayerCache(past_key_values, repeat_idx * self.config.num_hidden_layers)
                     if (past_key_values is not None and repeat_idx > 0)
@@ -405,7 +401,6 @@ class NandiModel(NandiPreTrainedModel):
                     position_ids=position_ids,
                     past_key_values=repeat_cache,
                     use_cache=use_cache,
-                    cache_position=cache_position,
                     **kwargs,
                 )