Fix UniversalTransformerCache.get_mask_sizes for batched generation (#8)

- Fix UniversalTransformerCache.get_mask_sizes for batched generation (7b9f90ad5eb766f3c3559124ad2fe67f885a2a5c)


Co-authored-by: Kristian Shw <KristianS7@users.noreply.huggingface.co>

modeling_ouro.py

@@ -195,6 +195,28 @@ class UniversalTransformerCache(Cache):
             return 0
         return cached.shape[2]
 
+    def get_mask_sizes(
+        self, cache_position: torch.Tensor, layer_idx: int = 0
+    ) -> tuple[int, int]:
+        """Return (kv_length, kv_offset) for attention mask creation.
+
+        The inherited ``Cache.get_mask_sizes`` falls back to
+        ``(cache_position.shape[0], 0)`` when ``layer_idx >= len(self.layers)``.
+        Because ``UniversalTransformerCache`` manages its own flat
+        ``key_cache`` / ``value_cache`` lists and keeps ``self.layers`` empty,
+        the fallback always fires.  During the prefill step this happens to be
+        correct (``cache_position`` spans the full input), but during
+        autoregressive decoding ``cache_position`` has length 1, so the mask is
+        built for ``kv_length=1`` instead of ``cached_length + 1``.  As a
+        result the 4D attention mask is too small and padding information is
+        lost, corrupting batched generation for every sequence except the
+        longest (unpadded) one.
+        """
+        query_length = cache_position.shape[0]
+        seq_length = self.get_seq_length(layer_idx)
+        kv_length = seq_length + query_length
+        return kv_length, 0
+
     def get_max_length(self) -> Optional[int]:
         return None