Fix CacheList batching (#769)
* Fix CacheList batching * Simplify fix * empty as method + change len to size --------- Co-authored-by: Awni Hannun <[email protected]>
This commit is contained in:
Tarjei Mandt
co-authored by
Awni Hannun
Awni Hannun
parent
5261ab85ee
commit
769069d66b
+3
-6
@@ -1003,7 +1003,8 @@ class BatchGenerator:
|
||||
self.uid_count += 1
|
||||
# Sort in ascending order of length
|
||||
self.unprocessed_prompts = sorted(
|
||||
self.unprocessed_prompts, key=lambda x: len(x[1]) + cache.cache_length(x[3])
|
||||
self.unprocessed_prompts,
|
||||
key=lambda x: len(x[1]) + max(c.size() for c in x[3]),
|
||||
)
|
||||
return uids
|
||||
|
||||
@@ -1023,10 +1024,6 @@ class BatchGenerator:
|
||||
|
||||
def _process_prompts(self, prompts):
|
||||
uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts)
|
||||
if hasattr(caches[0][0], "keys"):
|
||||
cache_is_empty = all(c[0].keys is None for c in caches)
|
||||
else:
|
||||
cache_is_empty = all(c[0][0] is None for c in caches)
|
||||
|
||||
lengths = [len(p) for p in inputs]
|
||||
max_length = max(lengths)
|
||||
@@ -1040,7 +1037,7 @@ class BatchGenerator:
|
||||
# New prompts so
|
||||
# 1. Left-pad the inputs
|
||||
# 2. Process
|
||||
if cache_is_empty:
|
||||
if all(c[0].empty() for c in caches):
|
||||
inputs = _left_pad_prompts(inputs, max_length=max_length)
|
||||
prompt_cache = _make_cache(self.model, padding)
|
||||
|
||||
|
||||
+53
-28
@@ -109,10 +109,6 @@ def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
|
||||
return [c.trim(num_tokens) for c in cache][0]
|
||||
|
||||
|
||||
def cache_length(cache: List[Any]):
|
||||
return max(len(c) for c in cache)
|
||||
|
||||
|
||||
def create_attention_mask(
|
||||
N: int, offset: int, return_array: bool, window_size: Optional[int]
|
||||
):
|
||||
@@ -146,23 +142,20 @@ class _BaseCache:
|
||||
def is_trimmable(self):
|
||||
return False
|
||||
|
||||
def __len__(self):
|
||||
"""The length of a cache is meant to represent the number of elements
|
||||
that we need to process in the attention. For instance for KVCache it
|
||||
is the size of the state, for RotatingKVCache it would be up to
|
||||
max_size etc."""
|
||||
def size(self):
|
||||
"""
|
||||
Return the size (i.e. sequence length) of the cache.
|
||||
|
||||
Not every cache is required to implement this, in which case the size
|
||||
will always be 0 (though the cache may not be empty).
|
||||
"""
|
||||
return 0
|
||||
|
||||
def __bool__(self):
|
||||
"""When an object defines __len__ then python defines the bool operator
|
||||
as len(obj) != 0. This, for instance, doesn't allow us to write
|
||||
|
||||
cache = cache or make_cache()
|
||||
|
||||
which is why we are overriding that behaviour with a constant bool
|
||||
operator return True.
|
||||
def empty(self):
|
||||
"""
|
||||
return True
|
||||
Return if the cache is empty or not.
|
||||
"""
|
||||
raise NotImplementedError("Cache sub-class must implement this.")
|
||||
|
||||
@classmethod
|
||||
def from_state(cls, state, meta_state):
|
||||
@@ -217,6 +210,9 @@ class ConcatenateKVCache(_BaseCache):
|
||||
def make_mask(self, *args, **kwargs):
|
||||
return create_attention_mask(*args, offset=self.offset, **kwargs)
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
|
||||
class QuantizedKVCache(_BaseCache):
|
||||
step = 256
|
||||
@@ -303,6 +299,9 @@ class QuantizedKVCache(_BaseCache):
|
||||
def make_mask(self, *args, **kwargs):
|
||||
return create_attention_mask(*args, offset=self.offset, **kwargs)
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
|
||||
class KVCache(_BaseCache):
|
||||
step = 256
|
||||
@@ -336,7 +335,7 @@ class KVCache(_BaseCache):
|
||||
self.values[..., prev : self.offset, :] = values
|
||||
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
|
||||
|
||||
def __len__(self):
|
||||
def size(self):
|
||||
return self.offset
|
||||
|
||||
@property
|
||||
@@ -379,6 +378,9 @@ class KVCache(_BaseCache):
|
||||
def merge(_, caches):
|
||||
return BatchKVCache.merge(caches)
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
|
||||
class RotatingKVCache(_BaseCache):
|
||||
step = 256
|
||||
@@ -487,7 +489,7 @@ class RotatingKVCache(_BaseCache):
|
||||
return self._update_in_place(keys, values)
|
||||
return self._update_concat(keys, values)
|
||||
|
||||
def __len__(self):
|
||||
def size(self):
|
||||
return min(self.offset, self.max_size)
|
||||
|
||||
@property
|
||||
@@ -554,6 +556,9 @@ class RotatingKVCache(_BaseCache):
|
||||
def merge(_, caches):
|
||||
return BatchRotatingKVCache.merge(caches)
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
|
||||
class ArraysCache(_BaseCache):
|
||||
def __init__(self, size, left_padding: Optional[List[int]] = None):
|
||||
@@ -631,6 +636,9 @@ class ArraysCache(_BaseCache):
|
||||
cache[e][i : i + 1] = caches[i][e]
|
||||
return cache
|
||||
|
||||
def empty(self):
|
||||
return self.cache[0] is None
|
||||
|
||||
|
||||
class MambaCache(ArraysCache):
|
||||
def __init__(self, left_padding: Optional[List[int]] = None):
|
||||
@@ -710,6 +718,9 @@ class ChunkedKVCache(_BaseCache):
|
||||
def meta_state(self, v):
|
||||
self.chunk_size, self.start_position = map(int, v)
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
|
||||
class CacheList(_BaseCache):
|
||||
def __init__(self, *caches):
|
||||
@@ -765,6 +776,20 @@ class CacheList(_BaseCache):
|
||||
def extract(self, idx):
|
||||
return CacheList(*(c.extract(idx) for c in self.caches))
|
||||
|
||||
def prepare(self, **kwargs):
|
||||
for c in self.caches:
|
||||
c.prepare(**kwargs)
|
||||
|
||||
def finalize(self):
|
||||
for c in self.caches:
|
||||
c.finalize()
|
||||
|
||||
def size(self):
|
||||
return max(c.size() for c in self.caches)
|
||||
|
||||
def empty(self):
|
||||
return self.caches[0].empty()
|
||||
|
||||
|
||||
def dynamic_roll(x, shifts, axis):
|
||||
n = x.shape[axis]
|
||||
@@ -830,9 +855,6 @@ class BatchKVCache(_BaseCache):
|
||||
self.values[..., prev : self._idx, :] = values
|
||||
return self.keys[..., : self._idx, :], self.values[..., : self._idx, :]
|
||||
|
||||
def __len__(self):
|
||||
return self._idx
|
||||
|
||||
def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
|
||||
if left_padding is not None:
|
||||
if self.keys is not None:
|
||||
@@ -938,7 +960,7 @@ class BatchKVCache(_BaseCache):
|
||||
|
||||
@classmethod
|
||||
def merge(cls, caches):
|
||||
lengths = [len(c) for c in caches]
|
||||
lengths = [c.size() for c in caches]
|
||||
max_length = max(lengths)
|
||||
padding = [max_length - l for l in lengths]
|
||||
B = len(caches)
|
||||
@@ -963,6 +985,9 @@ class BatchKVCache(_BaseCache):
|
||||
|
||||
return cache
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
|
||||
class BatchRotatingKVCache(_BaseCache):
|
||||
step = 256
|
||||
@@ -1096,9 +1121,6 @@ class BatchRotatingKVCache(_BaseCache):
|
||||
return self._update_in_place(keys, values)
|
||||
return self._update_concat(keys, values)
|
||||
|
||||
def __len__(self):
|
||||
return min(self._offset, self.max_size)
|
||||
|
||||
def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
|
||||
if left_padding is not None:
|
||||
if self.keys is not None:
|
||||
@@ -1252,7 +1274,7 @@ class BatchRotatingKVCache(_BaseCache):
|
||||
)
|
||||
|
||||
offsets = [c.offset for c in caches]
|
||||
lengths = [len(c) for c in caches]
|
||||
lengths = [c.size() for c in caches]
|
||||
max_length = max(lengths)
|
||||
padding = [max_length - l for l in lengths]
|
||||
B = len(caches)
|
||||
@@ -1277,3 +1299,6 @@ class BatchRotatingKVCache(_BaseCache):
|
||||
cache._offset = keys.shape[2]
|
||||
|
||||
return cache
|
||||
|
||||
def empty(self):
|
||||
return self.keys is None
|
||||
|
||||
Reference in New Issue
Block a user