cache.py

from __future__ import annotations

import torch
from exllamav2.ext import exllamav2_ext as ext_c

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from exllamav2.model import ExLlamaV2


class ExLlamaV2CacheBase:

    model = None
    max_seq_len: int
    batch_size: int

    current_seq_len: int

    key_states: list[torch.Tensor | None]
    value_states: list[torch.Tensor | None]
    num_key_value_heads: int
    num_hidden_layers: int
    head_dim: int

    dtype: torch.dtype
    weights_per_element: int
    has_scales: bool


    def __init__(self,
                 model: ExLlamaV2,
                 batch_size: int,
                 max_seq_len: int,
                 dtype: torch.dtype,
                 weights_per_element: int,
                 has_scales: bool):

        self.model = model
        self.max_seq_len = max_seq_len if max_seq_len != -1 else self.model.config.max_seq_len
        self.batch_size = batch_size
        self.dtype = dtype
        self.weights_per_element = weights_per_element
        self.has_scales = has_scales

        self.key_states = []
        self.value_states = []
        self.key_scales = []
        self.value_scales = []

        self.num_key_value_heads = self.model.config.num_key_value_heads
        self.num_hidden_layers = self.model.config.num_hidden_layers
        self.head_dim = self.model.config.head_dim

        self.current_seq_len = 0
        self.shape_basic = (self.batch_size, self.max_seq_len, self.num_key_value_heads, self.head_dim)
        self.shape_w = (self.batch_size, self.max_seq_len, self.num_key_value_heads, self.head_dim // self.weights_per_element)
        self.shape_s = (self.batch_size, self.max_seq_len, self.num_key_value_heads, self.head_dim // 32)


    def create_state_tensors(self,
                             copy_from: ExLlamaV2CacheBase | None,
                             lazy = False):

        assert copy_from is None or lazy == False, "Cannot use lazy cache initialization while copying"

        if copy_from:
            self.current_seq_len = copy_from.current_seq_len

        if not lazy:

            for i in range(self.num_hidden_layers):

                if copy_from is None:
                    device = self.model.cache_map[i]
                    p_key_states = torch.zeros(self.shape_w, dtype = self.dtype, device = device).contiguous()
                    p_value_states = torch.zeros(self.shape_w, dtype = self.dtype, device = device).contiguous()
                    if self.has_scales:
                        p_key_scales = torch.zeros(self.shape_s, dtype = torch.float16, device = device).contiguous()
                        p_value_scales = torch.zeros(self.shape_s, dtype = torch.float16, device = device).contiguous()
                else:
                    p_key_states = copy_from.key_states[i].clone()
                    p_value_states = copy_from.value_states[i].clone()
                    if self.has_scales:
                        p_key_scales = copy_from.key_scales[i].clone()
                        p_value_scales = copy_from.value_scales[i].clone()

                self.key_states.append(p_key_states)
                self.value_states.append(p_value_states)
                if self.has_scales:
                    self.key_scales.append(p_key_scales)
                    self.value_scales.append(p_value_scales)

        else:

            for i in range(self.num_hidden_layers):

                self.key_states.append(None)
                self.value_states.append(None)
                if self.has_scales:
                    self.key_scales.append(None)
                    self.value_scales.append(None)


    def update_cache_tensors(self):

        for k, v in self.model.cache_map.items():

            self.touch_device(v)

            if self.key_states[k] is not None:

                if str(self.key_states[k].device) == v: continue
                self.key_states[k] = None
                self.value_states[k] = None

            p_key_states = torch.zeros(self.shape_w, dtype = self.dtype, device = v).contiguous()
            p_value_states = torch.zeros(self.shape_w, dtype = self.dtype, device = v).contiguous()
            self.key_states[k] = p_key_states
            self.value_states[k] = p_value_states
            if self.has_scales:
                p_key_scales = torch.zeros(self.shape_s, dtype = torch.float16, device = v).contiguous()
                p_value_scales = torch.zeros(self.shape_s, dtype = torch.float16, device = v).contiguous()
                self.key_scales[k] = p_key_scales
                self.value_scales[k] = p_value_scales


    def roll_left(self):

        for i in range(self.model.config.num_hidden_layers):

            self.key_states[i] = torch.roll(self.key_states[i], shifts = -1, dims = 2)
            self.value_states[i] = torch.roll(self.value_states[i], shifts = -1, dims = 2)

        self.current_seq_len -= 1


    def get_kv_state(self, layer_idx: int, batch_size: int, offset: int, width: int) -> (torch.Tensor, torch.Tensor):
        raise NotImplementedError


    def store_kv_state(self, layer_idx: int, batch_size: int, offset: int, width: int):
        raise NotImplementedError


    def copy_states(self,
                    target: ExLlamaV2CacheBase,
                    from_column: int,
                    from_columns: int,
                    to_column: int,
                    to_columns: int,
                    from_row: int,
                    from_rows: int,
                    to_row: int,
                    to_rows: int):

        assert from_rows == 1
        assert from_columns == to_columns
        assert to_column + to_columns <= target.max_seq_len
        assert from_column + from_columns <= self.max_seq_len

        num_hidden_layers = self.model.config.num_hidden_layers

        for i in range(num_hidden_layers):

            source_view_k = self.key_states[i].narrow(0, from_row, from_rows).narrow(2, from_column, from_columns)
            source_view_v = self.value_states[i].narrow(0, from_row, from_rows).narrow(2, from_column, from_columns)
            target_view_k = target.key_states[i].narrow(0, to_row, to_rows).narrow(2, to_column, to_columns)
            target_view_v = target.value_states[i].narrow(0, to_row, to_rows).narrow(2, to_column, to_columns)

            if to_rows > 1:

                source_view_k = source_view_k.expand_as(target_view_k)
                source_view_v = source_view_v.expand_as(target_view_v)

            target_view_k.copy_(source_view_k)
            target_view_v.copy_(source_view_v)


    def touch_device(self, device):
        pass


class ExLlamaV2Cache(ExLlamaV2CacheBase):
    """
    FP16 cache
    """

    def __init__(self,
                 model: ExLlamaV2,
                 batch_size: int = 1,
                 max_seq_len: int = -1,
                 copy_from: ExLlamaV2Cache | None = None,
                 lazy: bool = False):

        super().__init__(model, batch_size, max_seq_len, torch.half, 1, False)

        self.create_state_tensors(copy_from, lazy)


    def get_kv_state(self,
                     layer_idx: int,
                     batch_size: int,
                     offset: int,
                     width: int) -> (torch.Tensor, torch.Tensor):

        return self.key_states[layer_idx], self.value_states[layer_idx]


    def store_kv_state(self,
                       layer_idx: int,
                       batch_size: int,
                       offset: int,
                       width: int):

        pass

    def dump_state(self,context_id,index,data_name,model_name):
        for i in range(self.num_hidden_layers):
            torch.save(self.key_states[i][:,:self.current_seq_len,:,:].clone(), f'./cacheDump/data/{data_name}/{model_name}/{index}_{context_id}_{i}_key.pt')
            torch.save(self.value_states[i][:,:self.current_seq_len,:,:].clone(), f'./cacheDump/data/{data_name}/{model_name}/{index}_{context_id}_{i}_value.pt') 
    @torch.inference_mode()
    def load_state(self,index,tokens_len):
        self.current_seq_len = 0
        system_len = tokens_len[0][0]
        query_len = tokens_len[0][-1]
        total_lens = 0
        for layer_index in range(self.num_hidden_layers):
            batch_keys, batch_values = self.get_kv_state(layer_index, 1, 0, 0)
            for context_id in range(0,12):
                if context_id == 0:
                    key = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_key.pt')[:,:system_len]
                    value = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_value.pt')[:,:system_len]
                    new_keys = batch_keys.narrow(1, 0, system_len)
                    new_values = batch_values.narrow(1, 0, system_len)
                    new_keys.copy_(key)
                    new_values.copy_(value)
                    total_lens += key.shape[1]
                elif context_id == 11:
                    # key = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_key.pt')[:,-query_len+1:]
                    # value = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_value.pt')[:,-query_len+1:]
                    # new_keys = batch_keys.narrow(1, sum(tokens_len[0][:context_id]), query_len-1)
                    # new_values = batch_values.narrow(1, sum(tokens_len[0][:context_id]), query_len-1)
                    # new_keys.copy_(key)
                    # new_values.copy_(value)
                    # total_lens += key.shape[1]
                    pass
                else:
                    key = torch.load(f'./cacheDump/data/{index}_{context_id}_{layer_index}_key.pt')[:,system_len:-query_len+1]
                    value  = torch.load(f'./cacheDump/data/{index}_{context_id}_{layer_index}_value.pt')[:,system_len:-query_len+1]
                    q_len = sum(tokens_len[0][:context_id])
                    new_keys = batch_keys.narrow(1, q_len, tokens_len[0][context_id])
                    new_values = batch_values.narrow(1, q_len, tokens_len[0][context_id])
                    new_keys.copy_(key)
                    new_values.copy_(value)
                    total_lens += key.shape[1]
        self.current_seq_len+=sum(tokens_len[0][:-1])
    def footprint(self):

        fp = []
        for layer in self.key_states + self.value_states:
            dev = layer.device.index
            while len(fp) <= dev: fp.append(0)
            fp[dev] += layer.numel() * 2
        return fp


    def clone(self):

        new = ExLlamaV2Cache(self.model, self.batch_size, self.max_seq_len, self)
        return new


class ExLlamaV2Cache_8bit(ExLlamaV2CacheBase):
    """
    8-bit cache. Keys and values are compressed to FP8 (e5m2) format by truncation.
    """

    def __init__(self,
                 model: ExLlamaV2,
                 batch_size: int = 1,
                 max_seq_len: int = -1,
                 copy_from: ExLlamaV2Cache_8bit | None = None,
                 lazy: bool = False):

        super().__init__(model, batch_size, max_seq_len, torch.uint8, 1, False)

        self.create_state_tensors(copy_from, lazy)

        # Create temp FP16 tensors for accessing FP8 layers

        self.temp_tensors = {}
        if not lazy:
            for device in self.model.get_cache_devices(): self.touch_device(device)


    def touch_device(self, device):

        if device in self.temp_tensors: return
        k = torch.zeros(self.shape_basic, dtype = torch.float16, device = device).contiguous()
        v = torch.zeros(self.shape_basic, dtype = torch.float16, device = device).contiguous()
        self.temp_tensors[device] = (k, v)


    def get_kv_state(self,
                     layer_idx: int,
                     batch_size: int,
                     offset: int,
                     width: int) -> (torch.Tensor, torch.Tensor):

        device = self.model.cache_map[layer_idx]
        temp_key_state, temp_value_state = self.temp_tensors[device]
        if width > 0: ext_c.fp8_to_fp16(self.key_states[layer_idx], temp_key_state, batch_size, offset, width)
        if width > 0: ext_c.fp8_to_fp16(self.value_states[layer_idx], temp_value_state, batch_size, offset, width)
        return temp_key_state, temp_value_state


    def store_kv_state(self,
                       layer_idx: int,
                       batch_size: int,
                       offset: int,
                       width: int):

        device = self.model.cache_map[layer_idx]
        temp_key_state, temp_value_state = self.temp_tensors[device]
        if width > 0: ext_c.fp16_to_fp8(temp_key_state, self.key_states[layer_idx], batch_size, offset, width)
        if width > 0: ext_c.fp16_to_fp8(temp_value_state, self.value_states[layer_idx], batch_size, offset, width)


    def footprint(self) -> list[int]:

        fp = []
        for layer in self.key_states + self.value_states:
            dev = layer.device.index
            while len(fp) <= dev: fp.append(0)
            fp[dev] += layer.numel() * 1
        for temp_k, temp_v in self.temp_tensors.values():
            fp[temp_k.device.index] += temp_k.numel() * 2
            fp[temp_v.device.index] += temp_v.numel() * 2
        return fp


    def clone(self) -> ExLlamaV2Cache_8bit:

        new = ExLlamaV2Cache_8bit(self.model, self.batch_size, self.max_seq_len, self)
        return new


class ExLlamaV2Cache_Q4(ExLlamaV2CacheBase):
    """
    Q4 cache. Uses grouped RTN quantization for keys/values
    """

    def __init__(self,
                 model: ExLlamaV2,
                 batch_size: int = 1,
                 max_seq_len: int = -1,
                 copy_from: ExLlamaV2Cache_Q4 | None = None,
                 lazy: bool = False):

        super().__init__(model, batch_size, max_seq_len, torch.uint8, 2, True)

        self.create_state_tensors(copy_from, lazy)

        # Create temp FP16 tensors for accessing Q4 layers

        self.temp_tensors = {}
        if not lazy:
            for device in self.model.get_cache_devices(): self.touch_device(device)


    def touch_device(self, device):

        if device in self.temp_tensors: return
        k = torch.zeros(self.shape_basic, dtype = torch.float16, device = device).contiguous()
        v = torch.zeros(self.shape_basic, dtype = torch.float16, device = device).contiguous()
        self.temp_tensors[device] = (k, v)


    def get_kv_state(self,
                     layer_idx: int,
                     batch_size: int,
                     offset: int,
                     width: int) -> (torch.Tensor, torch.Tensor):

        device = self.model.cache_map[layer_idx]
        temp_key_state, temp_value_state = self.temp_tensors[device]
        if width > 0: ext_c.q4_to_fp16_kv(self.key_states[layer_idx],
                                          temp_key_state,
                                          self.key_scales[layer_idx],
                                          self.value_states[layer_idx],
                                          temp_value_state,
                                          self.value_scales[layer_idx],
                                          batch_size,
                                          offset,
                                          width)
        return temp_key_state, temp_value_state


    def store_kv_state(self,
                       layer_idx: int,
                       batch_size: int,
                       offset: int,
                       width: int):

        device = self.model.cache_map[layer_idx]
        temp_key_state, temp_value_state = self.temp_tensors[device]
        if width > 0: ext_c.fp16_to_q4_kv(temp_key_state,
                                          self.key_states[layer_idx],
                                          self.key_scales[layer_idx],
                                          temp_value_state,
                                          self.value_states[layer_idx],
                                          self.value_scales[layer_idx],
                                          batch_size,
                                          offset,
                                          width)


    def footprint(self) -> list[int]:

        fp = []
        for layer in self.key_states + self.value_states:
            dev = layer.device.index
            while len(fp) <= dev: fp.append(0)
            fp[dev] += layer.numel() * 1
        for layer in self.key_scales + self.value_scales:
            dev = layer.device.index
            while len(fp) <= dev: fp.append(0)
            fp[dev] += layer.numel() * 2
        for temp_k, temp_v in self.temp_tensors.values():
            fp[temp_k.device.index] += temp_k.numel() * 2
            fp[temp_v.device.index] += temp_v.numel() * 2
        return fp


    def clone(self) -> ExLlamaV2Cache_Q4:

        new = ExLlamaV2Cache_Q4(self.model, self.batch_size, self.max_seq_len, self)
        return new

class ExLlamaV2Cache_CPU(ExLlamaV2CacheBase):
    """
    FP16 cache
    """

    def __init__(self,
                 model: ExLlamaV2,
                 batch_size: int = 1,
                 max_seq_len: int = -1,
                 copy_from: ExLlamaV2Cache_CPU | None = None,
                 lazy: bool = False):

        super().__init__(model, batch_size, max_seq_len, torch.half, 1, False)
        self.copy_stream = torch.cuda.Stream()
        self.create_state_tensors(copy_from, lazy)

    def create_state_tensors(self,
                             copy_from: ExLlamaV2CacheBase | None,
                             lazy = False):

        assert copy_from is None and lazy == False, "only supprot lazy == False"
        if not lazy:
            for i in range(self.num_hidden_layers):
                # device = self.model.cache_map[i]
                device = 'cpu'
                if copy_from is None:
                    p_key_states = torch.zeros(self.shape_w, dtype = self.dtype, device = device).contiguous()
                    p_value_states = torch.zeros(self.shape_w, dtype = self.dtype, device = device).contiguous()
                self.key_states.append(p_key_states)
                self.value_states.append(p_value_states)



    def get_kv_state(self,
                     layer_idx: int,
                     batch_size: int,
                     offset: int,
                     width: int) -> (torch.Tensor, torch.Tensor):
        self.copy_stream.synchronize()
        torch.cuda.current_stream().wait_stream(self.copy_stream)
        return self.key_states[layer_idx].to(self.model.cache_map[layer_idx]), self.value_states[layer_idx].to(self.model.cache_map[layer_idx])


    
    def store_kv_state(self, layer_idx , batch_size, all_len, key_states, value_states):
        ss = torch.cuda.current_stream()
        with torch.cuda.stream(self.copy_stream):
            torch.cuda.current_stream().wait_stream(ss)
            self.key_states[layer_idx][:batch_size,0:all_len,:,:] = key_states.to('cpu')
            self.value_states[layer_idx][:batch_size,0:all_len,:,:] = value_states.to('cpu')

    def footprint(self):

        fp = []
        for layer in self.key_states + self.value_states:
            dev = layer.device.index
            while len(fp) <= dev: fp.append(0)
            fp[dev] += layer.numel() * 2
        return fp


    def clone(self):

        new = ExLlamaV2Cache(self.model, self.batch_size, self.max_seq_len, self)
        return new
    

class ExLlamaV2Cache_importance(ExLlamaV2CacheBase):
    """
    FP16 cache
    """
    def __init__(self,
                 model: ExLlamaV2,
                 batch_size: int = 1,
                 max_seq_len: int = -1,
                 copy_from: ExLlamaV2Cache | None = None,
                 lazy: bool = False):

        super().__init__(model, batch_size, max_seq_len, torch.half, 1, False)
        self.shape_importance = (model.config.num_hidden_layers, self.max_seq_len, model.config.num_attention_heads)
        self.importance = torch.zeros(self.shape_importance)
        self.create_state_tensors(copy_from, lazy)


    def get_kv_state(self,
                     layer_idx: int,
                     batch_size: int,
                     offset: int,
                     width: int) -> (torch.Tensor, torch.Tensor):

        return self.key_states[layer_idx], self.value_states[layer_idx]


    def store_kv_state(self,
                       layer_idx: int,
                       batch_size: int,
                       offset: int,
                       width: int):

        pass

    def dump_state(self,context_id,index,data_name,model_name):
        for i in range(self.num_hidden_layers):
            torch.save(self.key_states[i][:,:self.current_seq_len,:,:].clone(), f'./cacheDump/data/{data_name}/{model_name}/{index}_{context_id}_{i}_key.pt')
            torch.save(self.value_states[i][:,:self.current_seq_len,:,:].clone(), f'./cacheDump/data/{data_name}/{model_name}/{index}_{context_id}_{i}_value.pt') 
    
    def dump_importance(self, batch, model_name, data_name, tokens_len):
        q_len = (self.importance[0].any(dim=-1)).sum()
        assert q_len == tokens_len
        for i in range(self.num_hidden_layers):
            torch.save(self.importance[i][:q_len,:].clone(), f'./cacheDump/data/{data_name}/{model_name}/{batch}_{i}_importance.pt')

    def load_importance(self, batch, model_name, data_name, tokens_len):
        for i in range(self.num_hidden_layers):
            self.importance[i][:tokens_len,:] = torch.load(f'./cacheDump/data/{data_name}/{model_name}/{batch}_{i}_importance.pt')
        
        return self.importance
    @torch.inference_mode()
    def load_state(self,index,tokens_len):
        self.current_seq_len = 0
        system_len = tokens_len[0][0]
        query_len = tokens_len[0][-1]
        total_lens = 0
        for layer_index in range(self.num_hidden_layers):
            batch_keys, batch_values = self.get_kv_state(layer_index, 1, 0, 0)
            for context_id in range(0,12):
                if context_id == 0:
                    key = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_key.pt')[:,:system_len]
                    value = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_value.pt')[:,:system_len]
                    new_keys = batch_keys.narrow(1, 0, system_len)
                    new_values = batch_values.narrow(1, 0, system_len)
                    new_keys.copy_(key)
                    new_values.copy_(value)
                    total_lens += key.shape[1]
                elif context_id == 11:
                    # key = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_key.pt')[:,-query_len+1:]
                    # value = torch.load(f'./cacheDump/data/{index}_{10}_{layer_index}_value.pt')[:,-query_len+1:]
                    # new_keys = batch_keys.narrow(1, sum(tokens_len[0][:context_id]), query_len-1)
                    # new_values = batch_values.narrow(1, sum(tokens_len[0][:context_id]), query_len-1)
                    # new_keys.copy_(key)
                    # new_values.copy_(value)
                    # total_lens += key.shape[1]
                    pass
                else:
                    key = torch.load(f'./cacheDump/data/{index}_{context_id}_{layer_index}_key.pt')[:,system_len:-query_len+1]
                    value  = torch.load(f'./cacheDump/data/{index}_{context_id}_{layer_index}_value.pt')[:,system_len:-query_len+1]
                    q_len = sum(tokens_len[0][:context_id])
                    new_keys = batch_keys.narrow(1, q_len, tokens_len[0][context_id])
                    new_values = batch_values.narrow(1, q_len, tokens_len[0][context_id])
                    new_keys.copy_(key)
                    new_values.copy_(value)
                    total_lens += key.shape[1]
        self.current_seq_len+=sum(tokens_len[0][:-1])
    def footprint(self):

        fp = []
        for layer in self.key_states + self.value_states:
            dev = layer.device.index
            while len(fp) <= dev: fp.append(0)
            fp[dev] += layer.numel() * 2
        return fp