Update modeling_alinlight.py

modeling_alinlight.py

@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# -*- coding: utf-8 -*-
+# Copyright 2026 EngineerGL Research.
+
 import math
 import torch
 import torch.nn as nn
@@ -20,7 +23,6 @@ import torch.nn.functional as F
 from typing import Optional, Tuple, List, Union
 from transformers import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
-
 from configuration_alinlight import AlinlightConfig
 
 class AlinlightRMSNorm(nn.Module):
@@ -44,11 +46,8 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.max_position_embeddings = max_position_embeddings
         self.scaling_factor = scaling_factor
         
-        # We calculate frequencies immediately upon initialization
         inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32).to(device) / self.dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        
-        # Initialize the cache
         self._set_cos_sin_cache(seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype())
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
@@ -60,10 +59,8 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
     def forward(self, x, seq_len=None):
-        # If the length is greater than the cache, we recalculate (a rare case, but needed for reliability)
         if seq_len > self.cos_cached.shape[0]:
              self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-        
         return (
             self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device),
             self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device)
@@ -139,14 +136,13 @@ class AlinlightAttention(nn.Module):
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
             value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
-        # === TPU OPTIMIZATION: Physical cropping instead of mask ===
+        # Truncation logic (TPU Optimization)
         if self.sliding_window is not None and key_states.shape[2] > self.sliding_window:
             key_states = key_states[:, :, -self.sliding_window:, :]
             value_states = value_states[:, :, -self.sliding_window:, :]
 
         past_key_value = (key_states, value_states) if use_cache else None
         
-        # GQA / MQA processing
         if self.num_key_value_groups > 1:
             key_states = key_states[:, :, None, :, :].expand(
                 bsz, self.num_key_value_heads, self.num_key_value_groups, key_states.shape[-2], self.head_dim
@@ -156,14 +152,18 @@ class AlinlightAttention(nn.Module):
                 bsz, self.num_key_value_heads, self.num_key_value_groups, value_states.shape[-2], self.head_dim
             ).reshape(bsz, self.num_heads, value_states.shape[-2], self.head_dim)
 
-        # SDPA (Flash Attention backend compatible)
+        # ✅ FIX: Динамический флаг is_causal
+        # Если q_len > 1 (обучение или prefill) -> True (маскируем будущее)
+        # Если q_len == 1 (генерация) -> False (видим всё прошлое, что есть в key_states)
+        is_causal = q_len > 1
+
         attn_output = F.scaled_dot_product_attention(
             query_states, 
             key_states, 
             value_states, 
             attn_mask=None, 
             dropout_p=0.0, 
-            is_causal=True 
+            is_causal=is_causal 
         )
         
         attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
@@ -189,7 +189,6 @@ class AlinlightDecoderLayer(nn.Module):
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
-        
         return hidden_states, None, present_key_value
 
 class AlinlightModel(PreTrainedModel):
@@ -221,7 +220,6 @@ class AlinlightModel(PreTrainedModel):
         else:
              inputs_embeds = kwargs.get("inputs_embeds")
 
-        # Length calculation for RoPE
         seq_len = inputs_embeds.shape[1]
         if past_key_values is not None:
              seq_len += past_key_values[0][0].shape[2]
@@ -264,7 +262,7 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
         super().__init__(config)
         self.model = AlinlightModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.lm_head.weight = self.model.embed_tokens.weight # Связываем веса (Weight Tying)
+        self.lm_head.weight = self.model.embed_tokens.weight
         self.post_init()
 
     def get_input_embeddings(self): return self.model.embed_tokens
@@ -273,7 +271,6 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
     def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-        # Optimization for generation: if there is a cache, we serve only the last token
         if past_key_values:
             input_ids = input_ids[:, -1:]