File size: 1,081 Bytes
2e30a77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import torch
import torch.nn as nn
class StaticKVCache(nn.Module):
def __init__(
self,
max_batch_size,
n_heads,
context_length,
head_dim,
device,
dtype=torch.float32
):
super().__init__()
self.register_buffer(
"k_cache",
torch.zeros(
(max_batch_size, n_heads, context_length, head_dim),
dtype=dtype,
device=device
)
)
self.register_buffer(
"v_cache",
torch.zeros(
(max_batch_size, n_heads, context_length, head_dim),
dtype=dtype,
device=device
)
)
def update(self, k_val, v_val, cache_position):
self.k_cache[:, :, cache_position] = k_val
self.v_cache[:, :, cache_position] = v_val
total_len = cache_position[-1].item() + 1
return (
self.k_cache[:, :, :total_len],
self.v_cache[:, :, :total_len]
) |