-
Notifications
You must be signed in to change notification settings - Fork 26.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Cache: Static cache as a standalone object #30476
Changes from 10 commits
fa0b05e
56cf73b
97aa785
a6853a1
326a1c7
9de9556
90771df
cd6aecd
b8caa5e
16e01f2
2d1eab8
bb159d2
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||||||
---|---|---|---|---|---|---|---|---|---|---|
|
@@ -61,6 +61,14 @@ def get_usable_length(self, new_seq_length: int, layer_idx: Optional[int] = 0) - | |||||||||
return max_length - new_seq_length | ||||||||||
return previous_seq_length | ||||||||||
|
||||||||||
def reorder_cache(self, beam_idx: torch.LongTensor): | ||||||||||
"""Reorders the cache for beam search, given the selected beam indices.""" | ||||||||||
for layer_idx in range(len(self.key_cache)): | ||||||||||
device = self.key_cache[layer_idx].device | ||||||||||
self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) | ||||||||||
device = self.value_cache[layer_idx].device | ||||||||||
self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) | ||||||||||
|
||||||||||
@property | ||||||||||
def seen_tokens(self): | ||||||||||
logger.warning_once( | ||||||||||
|
@@ -158,14 +166,6 @@ def get_max_length(self) -> Optional[int]: | |||||||||
"""Returns the maximum sequence length of the cached states. DynamicCache does not have a maximum length.""" | ||||||||||
return None | ||||||||||
|
||||||||||
def reorder_cache(self, beam_idx: torch.LongTensor): | ||||||||||
"""Reorders the cache for beam search, given the selected beam indices.""" | ||||||||||
for layer_idx in range(len(self.key_cache)): | ||||||||||
device = self.key_cache[layer_idx].device | ||||||||||
self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) | ||||||||||
device = self.value_cache[layer_idx].device | ||||||||||
self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) | ||||||||||
|
||||||||||
def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]: | ||||||||||
"""Converts the `DynamicCache` instance into the its equivalent in the legacy cache format.""" | ||||||||||
legacy_cache = () | ||||||||||
|
@@ -332,23 +332,14 @@ def update( | |||||||||
|
||||||||||
return self.key_cache[layer_idx], self.value_cache[layer_idx] | ||||||||||
|
||||||||||
def reorder_cache(self, beam_idx: torch.LongTensor): | ||||||||||
"""Reorders the cache for beam search, given the selected beam indices.""" | ||||||||||
for layer_idx in range(len(self.key_cache)): | ||||||||||
device = self.key_cache[layer_idx].device | ||||||||||
self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) | ||||||||||
device = self.value_cache[layer_idx].device | ||||||||||
self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) | ||||||||||
|
||||||||||
|
||||||||||
class StaticCache(Cache): | ||||||||||
""" | ||||||||||
Static Cache class to be used with `torch.compile(model)`. | ||||||||||
|
||||||||||
Parameters: | ||||||||||
config (`PretrainedConfig): | ||||||||||
The configuration file defining the `max_position_embeddings`, `hidden_size` and `num_attention_heads` | ||||||||||
required to initialize the static cache. | ||||||||||
The configuration file defining the shape-related attributes required to initialize the static cache. | ||||||||||
max_batch_size (`int`): | ||||||||||
The maximum batch size with which the model will be used. | ||||||||||
max_cache_len (`int`): | ||||||||||
|
@@ -373,9 +364,18 @@ def __init__(self, config: PretrainedConfig, max_batch_size: int, max_cache_len: | |||||||||
config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads | ||||||||||
) | ||||||||||
|
||||||||||
self.key_cache: List[torch.Tensor] = [] | ||||||||||
self.value_cache: List[torch.Tensor] = [] | ||||||||||
cache_shape = (max_batch_size, self.num_key_value_heads, self.max_cache_len, self.head_dim) | ||||||||||
self.key_cache: torch.Tensor = torch.zeros(cache_shape, dtype=self.dtype, device=device) | ||||||||||
self.value_cache: torch.Tensor = torch.zeros(cache_shape, dtype=self.dtype, device=device) | ||||||||||
for _ in range(config.num_hidden_layers): | ||||||||||
# Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph | ||||||||||
# breaks when updating the cache. | ||||||||||
new_layer_key_cache = torch.zeros(cache_shape, dtype=self.dtype, device=device) | ||||||||||
new_layer_value_cache = torch.zeros(cache_shape, dtype=self.dtype, device=device) | ||||||||||
torch._dynamo.mark_static_address(new_layer_key_cache) | ||||||||||
torch._dynamo.mark_static_address(new_layer_value_cache) | ||||||||||
self.key_cache.append(new_layer_key_cache) | ||||||||||
self.value_cache.append(new_layer_value_cache) | ||||||||||
|
||||||||||
def update( | ||||||||||
self, | ||||||||||
|
@@ -394,42 +394,38 @@ def update( | |||||||||
value_states (`torch.Tensor`): | ||||||||||
The new value states to cache. | ||||||||||
layer_idx (`int`): | ||||||||||
The index of the layer to cache the states for. Kept for backward compatibility | ||||||||||
The index of the layer to cache the states for. | ||||||||||
cache_kwargs (`Dict[str, Any]`, `optional`): | ||||||||||
Additional arguments for the cache subclass. The `StaticCache` just needs the `q_len` | ||||||||||
to know how much of the cache it should overwrite. | ||||||||||
Additional arguments for the cache subclass. The `StaticCache` needs the `cache_position` input | ||||||||||
to know how where to write in the cache. | ||||||||||
|
||||||||||
Return: | ||||||||||
A tuple containing the updated key and value states. | ||||||||||
""" | ||||||||||
new_cache_positions = cache_kwargs.get("cache_position") | ||||||||||
k_out = self.key_cache | ||||||||||
v_out = self.value_cache | ||||||||||
cache_position = cache_kwargs.get("cache_position") | ||||||||||
k_out = self.key_cache[layer_idx] | ||||||||||
v_out = self.value_cache[layer_idx] | ||||||||||
|
||||||||||
k_out[:, :, new_cache_positions] = key_states | ||||||||||
v_out[:, :, new_cache_positions] = value_states | ||||||||||
k_out[:, :, cache_position] = key_states | ||||||||||
v_out[:, :, cache_position] = value_states | ||||||||||
|
||||||||||
return k_out, v_out | ||||||||||
|
||||||||||
def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: | ||||||||||
"""Returns the sequence length of the cached states that were seen by the model. `layer_idx` kept for BC""" | ||||||||||
"""Returns the sequence length of the cached states that were seen by the model.""" | ||||||||||
# Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's | ||||||||||
# limit the check to the first batch member and head dimension. | ||||||||||
# TODO: This is error prone, a filled cache may be `0.0`. Let's use a stateless integer instead, after | ||||||||||
# https://github.com/pytorch/pytorch/issues/120248 is fixed | ||||||||||
return (self.key_cache[0, 0].any(dim=-1)).sum() | ||||||||||
return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() | ||||||||||
|
||||||||||
def get_max_length(self) -> Optional[int]: | ||||||||||
"""Returns the maximum sequence length of the cached states. DynamicCache does not have a maximum length.""" | ||||||||||
"""Returns the maximum sequence length of the cached states.""" | ||||||||||
return self.max_cache_len | ||||||||||
|
||||||||||
def reorder_cache(self, beam_idx: torch.LongTensor): | ||||||||||
"""Reorders the cache for beam search, given the selected beam indices.""" | ||||||||||
device = self.key_cache.device | ||||||||||
self.key_cache = self.key_cache.index_select(0, beam_idx.to(device)) | ||||||||||
device = self.value_cache.device | ||||||||||
self.value_cache = self.value_cache.index_select(0, beam_idx.to(device)) | ||||||||||
|
||||||||||
def to_legacy_cache(self): | ||||||||||
"""Dummy function for BC. We have to keep it because otherwise the call in the forward of models will break it""" | ||||||||||
return None | ||||||||||
def reset(self): | ||||||||||
"""Resets the cache values while preserving the objects""" | ||||||||||
for layer_idx in range(len(self.key_cache)): | ||||||||||
# In-place ops prevent breaking the static address | ||||||||||
self.key_cache[layer_idx] *= 0.0 | ||||||||||
self.value_cache[layer_idx] *= 0.0 | ||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
might be faster? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. setting to a new tensor produces a graph break 💔 (I'm assuming you meant There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No no, I think just filling them with zeros should work There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. That would result in But filling with zeros with There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ok 👍🏻 let's go with that then! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
will remove this one
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it's slow and not reliable, generate should never use it
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
(needs deprecation cycle and it's easer to do after we isolate the prefill stage, I'm going to leave it off this PR)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
fine by me to deprecate