Add memory analyzer & utomatically configure KV cache size (#6)

2023-03-11 23:23:14 -08:00
parent 1a7eb7da61
commit e9d3f2ff77
7 changed files with 216 additions and 34 deletions
--- a/cacheflow/models/init.py
+++ b/cacheflow/models/init.py
@@ -1,10 +1,12 @@
 from cacheflow.models.input_metadata import InputMetadata
+from cacheflow.models.model_utils import get_memory_analyzer
 from cacheflow.models.model_utils import get_model
-from cacheflow.models.model_utils import set_seed
+from cacheflow.models.utils import set_seed


 __all__ = [
    'InputMetadata',
+    'get_memory_analyzer',
    'get_model',
-    'set_seed'
+    'set_seed',
 ]