Fix UniversalTransformerCache.get_mask_sizes for batched generation

modeling_ouro.py

@@ -214,6 +214,18 @@ class UniversalTransformerCache(Cache):
             self.key_cache[idx] = key_entry.index_select(0, beam_idx.to(device))
             self.value_cache[idx] = value_entry.index_select(0, beam_idx.to(device))
 
+    def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int = 0) -> tuple[int, int]:
+        """Return (kv_length, kv_offset) accounting for cached tokens.
+
+        The inherited Cache.get_mask_sizes checks ``self.layers`` which is
+        always empty for UniversalTransformerCache, causing it to return
+        ``(query_length, 0)`` instead of ``(cached_length + query_length, 0)``
+        during autoregressive decoding.
+        """
+        query_length = cache_position.shape[0]
+        seq_length = self.get_seq_length(layer_idx)
+        return seq_length + query_length, 0
+
     @property
     def is_compileable(self) -> bool:
         return False