Update modeling_grok2.py

modeling_grok2.py

@@ -28,6 +28,7 @@ Architecture:
   Sparse MoE: 8 experts, top-2, SwiGLU (w1=gate, w3=up, w2=down)
   4x RMSNorm per layer (no bias)
   RoPE with scaled theta
+  KV cache disabled — forward pass only, no past_key_values
 """
 
 import math
@@ -157,7 +158,7 @@ class Grok2Attention(nn.Module):
         self.o_proj = nn.Linear(config.num_attention_heads * config.head_dim, config.hidden_size, bias=False)
         self.rotary_emb = Grok2RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta)
 
-    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
+    def forward(self, hidden_states, attention_mask=None, **kwargs):
         B, T, _ = hidden_states.shape
 
         q = self.q_proj(hidden_states).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
@@ -169,21 +170,6 @@ class Grok2Attention(nn.Module):
         sin = sin[:, :, :T, :self.head_dim]
         q, k = apply_rotary_emb(q, k, cos, sin)
 
-        if past_key_value is not None:
-            if hasattr(past_key_value, 'update'):
-                # DynamicCache — use the official update method
-                k, v = past_key_value.update(k, v, self._layer_idx)
-            elif hasattr(past_key_value, 'key_cache'):
-                layer_idx = getattr(self, '_layer_idx', 0)
-                if layer_idx < len(past_key_value.key_cache):
-                    k = torch.cat([past_key_value.key_cache[layer_idx], k], dim=2)
-                    v = torch.cat([past_key_value.value_cache[layer_idx], v], dim=2)
-            else:
-                k = torch.cat([past_key_value[0], k], dim=2)
-                v = torch.cat([past_key_value[1], v], dim=2)
-
-        present = (k, v) if use_cache else None
-
         # GQA expand
         k = k.repeat_interleave(self.num_kv_groups, dim=1)
         v = v.repeat_interleave(self.num_kv_groups, dim=1)
@@ -191,16 +177,14 @@ class Grok2Attention(nn.Module):
         scale = math.sqrt(self.head_dim)
         attn = torch.matmul(q, k.transpose(-2, -1)) / scale
 
-        # Attn logit softcapping
         if self.attn_softcap > 0:
             attn = attn / self.attn_softcap
             attn = torch.tanh(attn)
             attn = attn * self.attn_softcap
 
-        kv_len = k.shape[2]
         causal = torch.triu(
-            torch.full((T, kv_len), float("-inf"), device=q.device, dtype=q.dtype),
-            diagonal=1 + kv_len - T
+            torch.full((T, T), float("-inf"), device=q.device, dtype=q.dtype),
+            diagonal=1
         )
         attn = attn + causal.unsqueeze(0).unsqueeze(0)
 
@@ -210,12 +194,11 @@ class Grok2Attention(nn.Module):
         attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
         out = torch.matmul(attn, v)
         out = out.transpose(1, 2).contiguous().view(B, T, -1)
-        return self.o_proj(out), present
+        return self.o_proj(out)
 
 
 # ── MoE Expert ────────────────────────────────────────────────────────────────
 class Grok2Expert(nn.Module):
-    """Single expert: SwiGLU with w1=gate, w3=up, w2=down."""
     def __init__(self, hidden_size, moe_intermediate_size):
         super().__init__()
         self.w1 = nn.Linear(hidden_size, moe_intermediate_size, bias=False)
@@ -244,9 +227,8 @@ class Grok2SparseMoE(nn.Module):
         B, T, H = x.shape
         x_flat = x.view(-1, H)
 
-        router_logits = self.gate(x_flat)  # [B*T, n_experts]
+        router_logits = self.gate(x_flat)
 
-        # Router softcapping
         if self.router_softcap > 0:
             router_logits = router_logits / self.router_softcap
             router_logits = torch.tanh(router_logits)
@@ -268,7 +250,7 @@ class Grok2SparseMoE(nn.Module):
         return out.view(B, T, H)
 
 
-# ── Dense MLP (residual path) ─────────────────────────────────────────────────
+# ── Dense MLP ─────────────────────────────────────────────────────────────────
 class Grok2MLP(nn.Module):
     def __init__(self, config: Grok2Config):
         super().__init__()
@@ -285,27 +267,23 @@ class Grok2DecoderLayer(nn.Module):
     def __init__(self, config: Grok2Config, layer_idx: int):
         super().__init__()
         self.layer_idx = layer_idx
-        self.pre_attn_norm  = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.self_attn      = Grok2Attention(config)
-        self.self_attn._layer_idx = layer_idx
-        self.post_attn_norm = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.pre_moe_norm   = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.pre_attn_norm    = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn        = Grok2Attention(config)
+        self.post_attn_norm   = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.pre_moe_norm     = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.block_sparse_moe = Grok2SparseMoE(config)
-        self.mlp            = Grok2MLP(config)
-        self.post_moe_norm  = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp              = Grok2MLP(config)
+        self.post_moe_norm    = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
 
-    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
-        # Attention block
+    def forward(self, hidden_states, attention_mask=None, **kwargs):
+        # Attention
         residual = hidden_states
         hidden_states = self.pre_attn_norm(hidden_states)
-        hidden_states, present = self.self_attn(
-            hidden_states, attention_mask=attention_mask,
-            past_key_value=past_key_value, use_cache=use_cache
-        )
+        hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask)
         hidden_states = self.post_attn_norm(hidden_states)
         hidden_states = residual + hidden_states
 
-        # MoE + dense residual block
+        # MoE + dense residual
         residual = hidden_states
         hidden_states = self.pre_moe_norm(hidden_states)
         moe_out = self.block_sparse_moe(hidden_states)
@@ -313,7 +291,7 @@ class Grok2DecoderLayer(nn.Module):
         hidden_states = self.post_moe_norm(moe_out + mlp_out)
         hidden_states = residual + hidden_states
 
-        return hidden_states, present
+        return hidden_states
 
 
 # ── Model ─────────────────────────────────────────────────────────────────────
@@ -327,25 +305,11 @@ class Grok2Model(nn.Module):
         ])
         self.norm = Grok2RMSNorm(config.hidden_size, config.rms_norm_eps)
 
-    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False):
+    def forward(self, input_ids, attention_mask=None, **kwargs):
         hidden_states = self.embed_tokens(input_ids) * self.embedding_multiplier_scale
-        presents = [] if use_cache else None
-
-        for i, layer in enumerate(self.layers):
-            pkv = None
-            if past_key_values is not None:
-                if hasattr(past_key_values, 'key_cache'):
-                    pkv = past_key_values
-                else:
-                    pkv = past_key_values[i] if i < len(past_key_values) else None
-            hidden_states, present = layer(
-                hidden_states, attention_mask=attention_mask,
-                past_key_value=pkv, use_cache=use_cache
-            )
-            if use_cache and present is not None:
-                presents.append(present)
-
-        return self.norm(hidden_states), presents
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask=attention_mask)
+        return self.norm(hidden_states)
 
 
 # ── CausalLM ──────────────────────────────────────────────────────────────────
@@ -356,8 +320,8 @@ class Grok1ForCausalLM(PreTrainedModel, GenerationMixin):
 
     def __init__(self, config: Grok2Config):
         super().__init__(config)
-        self.model    = Grok2Model(config)
-        self.lm_head  = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.model   = Grok2Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.output_multiplier_scale = config.output_multiplier_scale
         self.final_logit_softcapping = config.final_logit_softcapping
         self.post_init()
@@ -372,13 +336,10 @@ class Grok1ForCausalLM(PreTrainedModel, GenerationMixin):
         past_key_values=None,
         inputs_embeds=None,
         labels=None,
-        use_cache=True,
+        use_cache=None,
         **kwargs,
     ):
-        hidden_states, presents = self.model(
-            input_ids, attention_mask=attention_mask,
-            past_key_values=past_key_values, use_cache=False
-        )
+        hidden_states = self.model(input_ids, attention_mask=attention_mask)
 
         logits = self.lm_head(hidden_states) * self.output_multiplier_scale
 
@@ -398,15 +359,15 @@ class Grok1ForCausalLM(PreTrainedModel, GenerationMixin):
             )
 
         return CausalLMOutputWithPast(
-            loss=loss, logits=logits, past_key_values=presents
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
         )
 
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-        if past_key_values is not None:
-            input_ids = input_ids[:, -1:]
-        return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
 
 
-# ── Register with AutoModel ───────────────────────────────────────────────────
+# ── Register ──────────────────────────────────────────────────────────────────
 AutoConfig.register("grok2", Grok2Config)
 AutoModelForCausalLM.register(Grok2Config, Grok1ForCausalLM)