22 lines
558 B
Python
22 lines
558 B
Python
import torch
|
|
|
|
|
|
def update_cache(fresh_indices, fresh_tokens, cache_dic, current, fresh_attn_map=None):
|
|
"""
|
|
Update the cache with the fresh tokens.
|
|
"""
|
|
step = current["step"]
|
|
layer = current["layer"]
|
|
module = current["module"]
|
|
# Update the cached tokens at the positions
|
|
|
|
indices = fresh_indices
|
|
|
|
cache_dic["cache"][-1][current["stream"]][current["layer"]][
|
|
current["module"]
|
|
].scatter_(
|
|
dim=1,
|
|
index=indices.unsqueeze(-1).expand(-1, -1, fresh_tokens.shape[-1]),
|
|
src=fresh_tokens,
|
|
)
|