Spaces:
Running on Zero
Running on Zero
File size: 3,655 Bytes
4a8de45 | 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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | from __future__ import annotations
from typing import Any
import torch
from torch import nn
class _LayerKV(nn.Module):
def __init__(self, k: torch.Tensor, v: torch.Tensor):
super().__init__()
self.register_buffer("k", k, persistent=False)
self.register_buffer("v", v, persistent=False)
class StaticShiftKVCache(nn.Module):
def __init__(
self,
num_layers: int,
window: int,
batch: int,
h_kv: int,
head_dim: int,
device,
dtype,
):
super().__init__()
self.window = int(window)
self.layers = nn.ModuleList(
[
_LayerKV(
k=torch.zeros(
(batch, h_kv, self.window, head_dim), device=device, dtype=dtype
),
v=torch.zeros(
(batch, h_kv, self.window, head_dim), device=device, dtype=dtype
),
)
for _ in range(num_layers)
]
)
self.is_sliding = [True] * num_layers
def get_seq_length(self) -> int:
return int(self.window)
def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int):
q_len = int(cache_position.numel()) if cache_position is not None else 1
if cache_position is None or cache_position.numel() == 0:
return self.window, 0
cur_pos = int(cache_position.reshape(-1)[-1].item())
kv_length = min(cur_pos + 1, self.window)
kv_offset = max(cur_pos + 1 - kv_length, 0)
return kv_length, kv_offset
def update(
self,
key_states: torch.Tensor,
value_states: torch.Tensor,
layer_idx: int,
cache_kwargs: Any,
):
layer = self.layers[layer_idx]
t_new = key_states.shape[-2]
if not torch._dynamo.is_compiling() and t_new > self.window:
raise ValueError(
f"StaticShiftKVCache assumes T_new <= window={self.window}, got {t_new}."
)
shift = min(int(t_new), self.window)
if shift < self.window:
layer.k[..., :-shift, :].copy_(layer.k[..., shift:, :])
layer.v[..., :-shift, :].copy_(layer.v[..., shift:, :])
layer.k[..., -shift:, :].copy_(key_states[..., -shift:, :])
layer.v[..., -shift:, :].copy_(value_states[..., -shift:, :])
else:
layer.k.copy_(key_states[..., -self.window :, :])
layer.v.copy_(value_states[..., -self.window :, :])
return layer.k, layer.v
def make_fixed_attention_mask(
batch: int,
window: int,
cache_position: torch.Tensor,
device: torch.device,
dtype: torch.dtype,
) -> torch.Tensor:
positions = cache_position.reshape(-1)
end_pos = positions[-1]
# Column j corresponds to the absolute key position:
# abs_key(j) = end_pos - window + 1 + j
# We keep the column index static so torch.compile sees a fixed-size buffer.
cols = torch.arange(window, device=device, dtype=positions.dtype).view(1, window)
min_valid_col = torch.clamp((window - 1) - end_pos, min=0)
max_valid_col = positions.view(-1, 1) - end_pos + (window - 1)
valid = (cols >= min_valid_col) & (cols <= max_valid_col)
zeros = torch.zeros((positions.numel(), window), device=device, dtype=dtype)
neg_inf = torch.full(
(positions.numel(), window), torch.finfo(dtype).min, device=device, dtype=dtype
)
mask = torch.where(valid, zeros, neg_inf)
return mask.view(1, 1, positions.numel(), window).expand(
batch, 1, positions.numel(), window
)
|