v4.46 support

variable_cache.py

@@ -32,18 +32,20 @@ class VariableCache(Cache_4_44_2, Cache):
     The cache of each layer is allocated to the same gpu as the layer itself.
     """
 
-    def __init__(self,
-                 config: DeciLMConfig,
-                 max_batch_size: int,
-                 max_cache_len: int | None,
-                 device: torch.device | str | None = None,
-                 dtype: torch.dtype | None = None,
-                 **kwargs: Any,
-                 ):
+    def __init__(
+            self,
+            config: DeciLMConfig,
+            batch_size: int = None,
+            max_cache_len: int = None,
+            device: torch.device = None,
+            dtype: torch.dtype = torch.float32,
+            max_batch_size: Optional[int] = None,
+            **kwargs: Any,
+    ) -> None:
         Cache_4_44_2.__init__(self)
 
         self.config = config
-        self.max_batch_size = max_batch_size
+        self.max_batch_size = batch_size or max_batch_size
         self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len
         self.dtype = dtype