MCPcopy Create free account
hub / github.com/Zefan-Cai/KVCache-Factory / batch_split

Method batch_split

pyramidkv/cache_utils_think.py:1396–1406  ·  view source on GitHub ↗

Split the current instance into a list of `DynamicCache` by the batch size. This will be used by `_split_model_inputs()` in `generation.utils`

(self, full_batch_size: int, split_size: int)

Source from the content-addressed store, hash-verified

1394 self.self_attention_cache.crop(maximum_length)
1395
1396 def batch_split(self, full_batch_size: int, split_size: int) -> "List[EncoderDecoderCache]":
1397 """Split the current instance into a list of `DynamicCache` by the batch size. This will be used by
1398 `_split_model_inputs()` in `generation.utils`"""
1399 self.check_dynamic_cache(self.batch_split.__name__)
1400 self_attention_cache = self.self_attention_cache.batch_split(full_batch_size, split_size)
1401 cross_attention_cache = self.cross_attention_cache.batch_split(full_batch_size, split_size)
1402
1403 out = []
1404 for self_attn, cross_attn in zip(self_attention_cache, cross_attention_cache):
1405 out.append(EncoderDecoderCache(self_attn, cross_attn))
1406 return out
1407
1408 @classmethod
1409 def from_batch_splits(cls, splits: List["EncoderDecoderCache"]) -> "EncoderDecoderCache":

Callers

nothing calls this directly

Calls 3

check_dynamic_cacheMethod · 0.95
EncoderDecoderCacheClass · 0.85
batch_splitMethod · 0.45

Tested by

no test coverage detected