Update veylon_model.py

veylon_model.py

@@ -384,11 +384,21 @@ class TransformerBlock(layers.Layer):
         return x + a + f
 
     def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
+        # Must mirror call()'s PARALLEL residual structure exactly:
+        # attn and ffn both read from norm(x) computed on the SAME
+        # pre-block x, then both get added to that same original x.
+        # Previously this was written as a sequential block (x=x+attn;
+        # x=x+ffn(norm2(x))) — ffn ended up conditioned on norm2(x+attn_out)
+        # instead of norm2(x), an input distribution the FFN weights were
+        # never trained on. That silently produced a different model at
+        # generation time than at training time: teacher-forced loss (via
+        # call()) never exercises this path and looked fine, while actual
+        # autoregressive generation compounded the mismatch every step.
         attn_out, nck, ncv = self.attn.generate_step(
             self.norm1(x), cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos,
         )
-        x = x + attn_out
-        x = x + self.ffn(self.norm2(x), training=False)
+        ffn_out = self.ffn(self.norm2(x), training=False)
+        x = x + attn_out + ffn_out
         return x, nck, ncv
 
     def get_config(self):