model.py

import sys
min_version = (3, 9)
if sys.version_info < min_version:
    print("")
    print(f" ## Warning: this project requires Python {min_version[0]}.{min_version[1]} or higher.")
    print("")

import torch
from torch import nn
import torch.nn.functional as F
from safetensors import safe_open
import cuda_ext
import json
import math
import gc
from enum import Enum

try:
    from flash_attn import flash_attn_func
except:
    pass

class ParsedEnum(Enum):

    def __str__(self):
        return self.name.lower()

    def __repr__(self):
        return str(self)

    @classmethod
    def argparse(cls, s):
        try:
            return cls[s.upper()]
        except KeyError:
            return s


class ExLlamaConfig:

    # Load config from Llama config.json

    def __init__(self, model_config_path):

        with open(model_config_path) as f:
            read_config = json.load(f)

        # Loaded/automatic settings

        self.bos_token_id = read_config["bos_token_id"] if "bos_token_id" in read_config else 1
        self.eos_token_id = read_config["eos_token_id"] if "eos_token_id" in read_config else 2
        self.pad_token_id = read_config["pad_token_id"] if "pad_token_id" in read_config else 0

        self.hidden_size = read_config["hidden_size"]
        self.initializer_range = read_config["initializer_range"]
        self.intermediate_size = read_config["intermediate_size"]
        self.num_attention_heads = read_config["num_attention_heads"]
        self.num_hidden_layers = read_config["num_hidden_layers"]
        self.rms_norm_eps = read_config["rms_norm_eps"]
        self.vocab_size = read_config["vocab_size"]

        if "num_key_value_heads" in read_config:
            self.num_key_value_heads = read_config["num_key_value_heads"]
            self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
        else:
            self.num_key_value_heads = self.num_attention_heads
            self.num_key_value_groups = 1

        self.rotary_embedding_base = read_config["rope_theta"] if "rope_theta" in read_config else 10000.0
        self.head_dim = self.hidden_size // self.num_attention_heads

        self.groupsize = None  # Autodetected
        self.act_order = False  # Autodetected
        self.empty_g_idx = False  # Autodetected

        # Required settings

        self.model_path = None  # str or list[str]
        self.device_map = ExLlamaDeviceMap(self.num_hidden_layers)

        # Optional settings

        self.max_seq_len = 2048  # Reduce to save memory. Can also be increased, ideally while also using compress_pos_emn and a compatible model/LoRA
        self.max_input_len = 2048  # Maximum length of input IDs in a single forward pass. Sequences longer than this will be processed in multiple steps
        self.max_attention_size = 2048**2  # Sequences will be processed in chunks to keep the size of the attention weights matrix <= this
        self.compress_pos_emb = 1.0  # Increase to compress positional embeddings applied to sequence
        self.alpha_value = 1.0 # Alpha value for NTK RoPE scaling. Similar to compress_pos_emb, higher values increaste ctx but add Perplexity.
        self.gpu_peer_fix = False # Apparently Torch can have problems transferring tensors directly one GPU to another sometimes. Enable this to expliticly move tensors via system RAM instead, where needed
        self.auto_map = None  # List of floats with memory allocation in GB, per CUDA device, overrides device_map

        # Tuning

        self.use_flash_attn_2 = False
        self.matmul_recons_thd = 8
        self.fused_mlp_thd = 2
        self.sdp_thd = 8
        self.fused_attn = True
        self.matmul_fused_remap = False
        self.rmsnorm_no_half2 = False
        self.rope_no_half2 = False
        self.matmul_no_half2 = False
        self.silu_no_half2 = False
        self.concurrent_streams = False

    # Copy tuning params to C++ extension

    def set_tuning_params(self):

        cuda_ext.exllama_ext.set_tuning_params(self.matmul_recons_thd,
                                               self.fused_mlp_thd,
                                               self.sdp_thd,
                                               self.matmul_fused_remap,
                                               self.rmsnorm_no_half2,
                                               self.rope_no_half2,
                                               self.matmul_no_half2,
                                               self.silu_no_half2,
                                               self.concurrent_streams)

    # Parse and set list of GPU VRAM allocations

    def set_auto_map(self, map_string):

        if map_string is None: self.auto_map = None
        else: self.auto_map = [float(alloc) for alloc in map_string.split(",")]

    def calculate_rotary_embedding_base(self):
        self.rotary_embedding_base = self.rotary_embedding_base * self.alpha_value ** (self.head_dim / (self.head_dim-2))


# 4-bit linear layer implementation

class Ex4bitLinear:

    def __init__(self, config, in_features, out_features, has_bias, tensors, key):

        self.config = config
        self.key = key
        self.in_features = in_features
        self.out_features = out_features

        self.qweight = tensors[key + ".qweight"]
        self.qzeros = tensors[key + ".qzeros"]
        self.scales = tensors[key + ".scales"]
        self.g_idx = tensors[key + ".g_idx"].cpu() if key + ".g_idx" in tensors else None
        self.bias = tensors[key + ".bias"] if has_bias else None

        if self.g_idx is not None and (self.g_idx == 0).all():
            self.config.empty_g_idx = True
            self.g_idx = None

        self.device = self.qweight.device
        self.device_index = self.device.index

        self.q4 = cuda_ext.ext_make_q4(self.qweight,
                                       self.qzeros,
                                       self.scales,
                                       self.g_idx,
                                       self.device_index)

        self.height = tensors[key + ".qweight"].shape[0] * 8
        self.width = tensors[key + ".qweight"].shape[1]

        # Infer groupsize from height of qzeros

        self.groupsize = None
        if self.qzeros.shape[0] > 1:
            self.groupsize = (self.qweight.shape[0] * 8) // self.qzeros.shape[0]
            if self.config.groupsize is None:
                self.config.groupsize = self.groupsize


        # Handle act-order matrix

        if self.g_idx is not None:

            if self.groupsize is None: raise ValueError("Found group index but no groupsize. What do?")
            self.config.act_order = True


    def lora_applies(self, lora):

        if lora is None: return False
        return self.key + ".lora_A.weight" in lora.tensors


    def lora_apply(self, lora, x):

        lora_a = lora.tensors[self.key + ".lora_A.weight"]
        lora_b = lora.tensors[self.key + ".lora_B.weight"]
        out = torch.matmul(x, lora_a)
        out = torch.matmul(out, lora_b)
        # out = cuda_ext.ext_half_matmul(x, lora_a.contiguous(), cublas = True)
        # out = cuda_ext.ext_half_matmul(out, lora_b.contiguous(), cublas = True)
        return out


    def get_lora_tensors_or_meta(self, lora):

        if not self.lora_applies(lora):
            return cuda_ext.none_tensor, cuda_ext.none_tensor
        else:
            lora_a = lora.tensors[self.key + ".lora_A.weight"]
            lora_b = lora.tensors[self.key + ".lora_B.weight"]
            return lora_a, lora_b


    def forward(self, x, lora):

        if self.lora_applies(lora):
            lora_a = lora.tensors[self.key + ".lora_A.weight"]
            lora_b = lora.tensors[self.key + ".lora_B.weight"]
            out = cuda_ext.ext_q4_matmul(x, self.q4, self.width, lora_a, lora_b)
        else:
            out = cuda_ext.ext_q4_matmul(x, self.q4, self.width)

        # out = cuda_ext.ext_q4_matmul(x, self.q4, self.width)
        # if self.lora_applies(lora):
        #     out += self.lora_apply(lora, x)

        if self.bias is not None: out.add_(self.bias)
        return out


# Llama MLP

class ExLlamaMLP:

    def __init__(self, config, tensors, key):

        self.config = config

        self.gate_proj = Ex4bitLinear(config, self.config.hidden_size, self.config.intermediate_size, False, tensors, key + ".gate_proj")
        self.up_proj = Ex4bitLinear(config, self.config.hidden_size, self.config.intermediate_size, False, tensors, key + ".up_proj")
        self.down_proj = Ex4bitLinear(config, self.config.intermediate_size, self.config.hidden_size, False, tensors, key + ".down_proj")

        self.act_fn = nn.SiLU()

    def fused(self, x, buffer, post_attention_layernorm, lora):

        bsz, q_len, _ = x.size()

        gate_a, gate_b = self.gate_proj.get_lora_tensors_or_meta(lora)
        up_a, up_b = self.up_proj.get_lora_tensors_or_meta(lora)
        down_a, down_b = self.down_proj.get_lora_tensors_or_meta(lora)

        temp_size = 0
        if not gate_a.is_meta: temp_size = max(temp_size, bsz * q_len * gate_a.shape[1])
        if not up_a.is_meta:   temp_size = max(temp_size, bsz * q_len * up_a.shape[1])
        if not down_a.is_meta: temp_size = max(temp_size, bsz * q_len * down_a.shape[1])

        if temp_size > 0: lora_temp = torch.empty((1, temp_size), dtype = torch.float16, device = x.device)
        else: lora_temp = cuda_ext.none_tensor

        cuda_ext.exllama_ext.q4_mlp(x.view(-1, x.shape[-1]),
                                    post_attention_layernorm.weight,
                                    self.config.rms_norm_eps,
                                    self.gate_proj.q4,
                                    self.up_proj.q4,
                                    self.down_proj.q4,
                                    gate_a, gate_b,
                                    up_a, up_b,
                                    down_a, down_b,
                                    lora_temp)


    def forward(self, x, buffer, lora):

        y = self.gate_proj.forward(x, lora)
        y = self.act_fn(y)
        y *= self.up_proj.forward(x, lora)
        y = self.down_proj.forward(y, lora)

        return y


# RMS Layer norm.

class ExLlamaRMSNorm:

    def __init__(self, config, tensors, key):

        self.config = config
        self.variance_epsilon = self.config.rms_norm_eps
        self.weight = tensors[key]


    def forward(self, hidden_states, buffer):

        hidden_states = cuda_ext.ext_rms_norm(hidden_states, self.weight, self.variance_epsilon)
        return hidden_states


# Llama attention

class ExLlamaAttention:

    def __init__(self, config, tensors, key, sin, cos, index):

        self.config = config
        self.sin = sin
        self.cos = cos
        self.index = index

        self.q_proj = Ex4bitLinear(config, self.config.hidden_size, self.config.num_attention_heads * self.config.head_dim, False, tensors, key + ".q_proj")
        self.k_proj = Ex4bitLinear(config, self.config.hidden_size, self.config.num_key_value_heads * self.config.head_dim, False, tensors, key + ".k_proj")
        self.v_proj = Ex4bitLinear(config, self.config.hidden_size, self.config.num_key_value_heads * self.config.head_dim, False, tensors, key + ".v_proj")
        self.o_proj = Ex4bitLinear(config, self.config.num_attention_heads * self.config.head_dim, self.config.hidden_size, False, tensors, key + ".o_proj")


    def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:

        # TODO: This seems inefficient. It should be possible to broadcast in the attention matmul to avoid building
        # temporary K/V tensors like this

        batch, num_key_value_heads, slen, head_dim = hidden_states.shape
        if n_rep == 1: return hidden_states

        hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
        return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


    def fused(self, hidden_states, cache, buffer, input_layernorm, lora):

        bsz, q_len, _ = hidden_states.size()
        past_len = cache.current_seq_len

        # Lora tensors

        q_a, q_b = self.q_proj.get_lora_tensors_or_meta(lora)
        k_a, k_b = self.k_proj.get_lora_tensors_or_meta(lora)
        v_a, v_b = self.v_proj.get_lora_tensors_or_meta(lora)
        o_a, o_b = self.o_proj.get_lora_tensors_or_meta(lora)

        temp_size = 0
        if not q_a.is_meta: temp_size = max(temp_size, bsz * q_len * q_a.shape[1])
        if not k_a.is_meta: temp_size = max(temp_size, bsz * q_len * k_a.shape[1])
        if not v_a.is_meta: temp_size = max(temp_size, bsz * q_len * v_a.shape[1])
        if not o_a.is_meta: temp_size = max(temp_size, bsz * q_len * o_a.shape[1])
        if temp_size > 0: lora_temp = torch.empty((1, temp_size), dtype = torch.float16, device = hidden_states.device)
        else: lora_temp = cuda_ext.none_tensor

        # Project q, k, v, apply position embeddings to k and v, update cache

        query_states = torch.empty((bsz, q_len, self.config.num_attention_heads * self.config.head_dim), dtype = torch.float16, device = hidden_states.device)
        key_states = torch.empty((bsz, q_len, self.config.num_key_value_heads * self.config.head_dim), dtype = torch.float16, device = hidden_states.device)
        value_states = torch.empty((bsz, q_len, self.config.num_key_value_heads * self.config.head_dim), dtype = torch.float16, device = hidden_states.device)

        cuda_ext.exllama_ext.q4_attn(hidden_states,
                                     input_layernorm.weight,
                                     self.config.rms_norm_eps,
                                     query_states,
                                     key_states,
                                     value_states,
                                     self.q_proj.q4,
                                     self.k_proj.q4,
                                     self.v_proj.q4,
                                     self.sin,
                                     self.cos,
                                     q_len,
                                     past_len,
                                     self.config.num_attention_heads,
                                     self.config.num_key_value_heads,
                                     self.config.head_dim,
                                     cache.key_states[self.index],
                                     cache.value_states[self.index],
                                     cache.max_seq_len,
                                     q_a, q_b,
                                     k_a, k_b,
                                     v_a, v_b,
                                     lora_temp)

        query_states = query_states.view(bsz, q_len, self.config.num_attention_heads, self.config.head_dim)

        # Get k, v with past

        key_states = cache.key_states[self.index].narrow(2, 0, past_len + q_len).narrow(0, 0, bsz)
        value_states = cache.value_states[self.index].narrow(2, 0, past_len + q_len).narrow(0, 0, bsz)

        # Repeat K/V heads if num_key_value_headsn_kv_heads < n_heads

        query_states.transpose_(1, 2)
        key_states = self.repeat_kv(key_states, self.config.num_key_value_groups)
        value_states = self.repeat_kv(value_states, self.config.num_key_value_groups)

        # Attention
        # TODO: Figure out if we can use cublasHgemmStridedBatched() to do this matmul without reshaping. Torch uses
        # gemmStridedBatchedEx() internally, so it should be possible.

        # -- Flash Attention 2.0

        if self.config.use_flash_attn_2 and (past_len == 0 or q_len == 1):

            key_states = key_states.transpose(1, 2)
            value_states = value_states.transpose(1, 2)
            query_states = query_states.transpose(1, 2)
            attn_output = flash_attn_func(query_states, key_states, value_states, causal = (past_len == 0))

        # -- HF Transformers regular attention, faster on shorter sequences, same VRAM usage

        else:

            key_states.transpose_(2, 3)
            attn_weights = torch.matmul(query_states, key_states)
            attn_weights /= math.sqrt(self.config.head_dim)
            attn_weights = nn.functional.softmax(attn_weights, dim = -1, dtype = torch.float16)
            attn_output = torch.matmul(attn_weights, value_states)
            attn_output = attn_output.transpose(1, 2)

        attn_output = attn_output.reshape(bsz, q_len, self.config.hidden_size)

        # Output projection

        cuda_ext.exllama_ext.q4_attn_2(hidden_states,
                                       attn_output,
                                       self.o_proj.q4,
                                       o_a, o_b,
                                       lora_temp)
        # return hidden_states


    def forward(self, hidden_states, cache, buffer, lora):

        bsz, q_len, _ = hidden_states.size()
        past_len = cache.current_seq_len

        # Project q, k, v, apply position embeddings to k and v

        query_states = self.q_proj.forward(hidden_states, lora)
        key_states = self.k_proj.forward(hidden_states, lora)

        cuda_ext.exllama_ext.rope_(query_states, self.sin, self.cos, past_len, self.config.num_attention_heads, self.config.head_dim)
        cuda_ext.exllama_ext.rope_(key_states, self.sin, self.cos, past_len, self.config.num_key_value_heads, self.config.head_dim)

        query_states = query_states.view(bsz, q_len, self.config.num_attention_heads, self.config.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)
        value_states = self.v_proj.forward(hidden_states, lora).view(bsz, q_len, self.config.num_key_value_heads, self.config.head_dim).transpose(1, 2)

        # Add keys and values to cache

        new_keys = cache.key_states[self.index].narrow(2, past_len, q_len).narrow(0, 0, bsz)
        new_values = cache.value_states[self.index].narrow(2, past_len, q_len).narrow(0, 0, bsz)
        new_keys.copy_(key_states)
        new_values.copy_(value_states)

        # Key/value tensors with past

        key_states = cache.key_states[self.index].narrow(2, 0, past_len + q_len).narrow(0, 0, bsz)
        value_states = cache.value_states[self.index].narrow(2, 0, past_len + q_len).narrow(0, 0, bsz)

        # Attention

        # -- Flash Attention 2.0

        if self.config.use_flash_attn_2 and (past_len == 0 or q_len == 1):

            key_states = key_states.transpose(1, 2)
            value_states = value_states.transpose(1, 2)
            query_states = query_states.transpose(1, 2)
            attn_output = flash_attn_func(query_states, key_states, value_states, causal = (past_len == 0))

        # -- HF Transformers regular attention, faster on shorter sequences, same VRAM usage

        elif self.config.sdp_thd == 0 or q_len < self.config.sdp_thd:

            key_states = self.repeat_kv(key_states, self.config.num_key_value_groups)
            value_states = self.repeat_kv(value_states, self.config.num_key_value_groups)

            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3))
            attn_weights /= math.sqrt(self.config.head_dim)
            if buffer.attn_mask is not None: attn_weights = attn_weights + buffer.attn_mask
            attn_weights = nn.functional.softmax(attn_weights, dim = -1, dtype = torch.float16)
            attn_output = torch.matmul(attn_weights, value_states)
            attn_output = attn_output.transpose(1, 2)

        # -- Scaled dot-product attention from PyTorch 2, should be comparable to xformers (?)

        else:

            # Torch's SDP attention has a built-in causal mask feature which we can use only when there is no past, i.e.
            # it can only apply a square attention mask. It saves quite a bit of VRAM but in practice Torch seems to use
            # the same amount of memory at peak anyway.
            #
            # TODO: Apparently flash attention is disabled when supplying an attention mask tensor. Figure out if this
            # is true and maybe drop SDP altogether. If causal masking in flash-attn is updated eventually there should
            # be no need for this anyway.

            key_states = self.repeat_kv(key_states, self.config.num_key_value_groups)
            value_states = self.repeat_kv(value_states, self.config.num_key_value_groups)

            if past_len > 0 or (bsz > 1 and buffer.attn_mask is not None):
                attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask = buffer.attn_mask, is_causal = False)
            else:
                attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask = None, is_causal = True)

            attn_output = attn_output.transpose(1, 2)

        # Output projection

        attn_output = attn_output.reshape(bsz, q_len, self.config.hidden_size)
        attn_output = self.o_proj.forward(attn_output, lora)

        return attn_output


def _rows(x):
    xdp = 1
    for y in x.shape[:-1]: xdp *= y
    return xdp

class ExLlamaDecoderLayer:

    def __init__(self, config, tensors, key, index, sin, cos):

        self.config = config
        self.index = index

        self.self_attn = ExLlamaAttention(self.config, tensors, key + ".self_attn", sin, cos, self.index)
        self.mlp = ExLlamaMLP(self.config, tensors, key + ".mlp")

        self.input_layernorm = ExLlamaRMSNorm(self.config, tensors, key + ".input_layernorm.weight")
        self.post_attention_layernorm = ExLlamaRMSNorm(self.config, tensors, key + ".post_attention_layernorm.weight")


    def forward(self, hidden_states, cache, buffer, lora):

        # Self-attention

        if self.config.fused_attn and _rows(hidden_states) == 1:

            self.self_attn.fused(hidden_states, cache, buffer, self.input_layernorm, lora)

        else:

            residual = hidden_states
            hidden_states = self.input_layernorm.forward(hidden_states, buffer)
            hidden_states = self.self_attn.forward(hidden_states, cache, buffer, lora)
            hidden_states = residual + hidden_states

        # MLP

        if self.config.fused_mlp_thd > 0 and _rows(hidden_states) <= self.config.fused_mlp_thd:

            self.mlp.fused(hidden_states, buffer, self.post_attention_layernorm, lora)

        else:

            residual = hidden_states
            hidden_states = self.post_attention_layernorm.forward(hidden_states, buffer)
            hidden_states = self.mlp.forward(hidden_states, buffer, lora)
            hidden_states = residual + hidden_states

        return hidden_states


# Persistent cache for inference. Allocate the whole thing up front.

class ExLlamaCache:

    def __init__(self, model, batch_size = 1, max_seq_len = -1, copy_from = None):

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

        self.key_states = []
        self.value_states = []
        self.current_seq_len = 0

        # Preallocate full-length cache

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

            if copy_from is None:

                p_key_states = torch.zeros(self.batch_size, self.config.num_key_value_heads, self.max_seq_len, self.config.head_dim, dtype = torch.float16, device = self.model.config.device_map.layers[i])
                p_value_states = torch.zeros(self.batch_size, self.config.num_key_value_heads, self.max_seq_len, self.config.head_dim, dtype = torch.float16, device = self.model.config.device_map.layers[i])

            else:

                p_key_states = copy_from.key_states[i].clone()
                p_value_states = copy_from.value_states[i].clone()

            self.key_states.append(p_key_states)
            self.value_states.append(p_value_states)


    def zero(self):

        for i in range(self.config.num_hidden_layers):
            self.key_states[i].zero_()
            self.value_states[i].zero_()


    def clone(self):

        new = ExLlamaCache(self.model, batch_size = self.batch_size, max_seq_len = self.max_seq_len, copy_from = self)
        return new


    def roll_left(self):

        for i in range(self.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 copy_states(self, target, from_column, from_columns, to_column, to_columns, from_row, from_rows, to_row, to_rows):

        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

        for i in range(self.config.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)


# Device map for the model.

class ExLlamaDeviceMap:

    def __init__(self, num_layers):

        self.num_layers = num_layers

        self.embed_tokens = "cpu"  # Embedding table on CPU saves 400 MB on the 30B model with no measurable impact on performance
        self.lm_head = "cuda:0"
        self.norm = "cuda:0"
        self.layers = ["cuda:0"] * self.num_layers


    def get_layers_devs(self):

        return sorted(list(set(self.layers)))


    def get_all_devs(self):

        return sorted(list(set(self.layers + [self.lm_head, self.norm, self.embed_tokens])))


    def map(self, key):

        if key.startswith("lm_head."): return self.lm_head
        if key.startswith("model.embed_tokens."): return self.embed_tokens
        if key.startswith("model.norm."): return self.norm

        if key.startswith("model.layers."):
            num = int(key.split(".")[2])
            return self.layers[num]

        raise ValueError("Unknown key: " + key)


class ExLlamaBuffer:

    config: ExLlamaConfig

    def __init__(self, config):

        self.config = config

    # Attention mask

    attn_mask: torch.Tensor = None

    # Move to device

    def to(self, device):

        new = ExLlamaBuffer(self.config)
        new.attn_mask = None if self.attn_mask is None else _move_tensor(self.attn_mask, device, "attn_mask", self.config)
        return new


def _device_to_int(device):

    return int(device[device.find(":") + 1:])

def _skip_key(key):

    if key.endswith("_proj.bias"): return True
    if key.endswith(".rotary_emb.inv_freq"): return True
    return False

def _move_tensor(tensor, new_device, name, config):
    device = str(tensor.device)
    if device == new_device: return tensor
    if config.gpu_peer_fix:
        if str(device).startswith("cuda:") and str(new_device).startswith("cuda:"):
            tensor = tensor.to("cpu")
    return tensor.to(new_device)

def _layer_dtype_size(key):
    if key.endswith(".weight"): return 2
    if key.endswith(".qweight"): return 4
    if key.endswith(".qzeros"): return 4
    if key.endswith(".scales"): return 2
    if key.endswith(".g_idx"): return 0
    raise ValueError("Unrecognized layer: " + key)


class ExLlama:

    def __init__(self, config):

        self.config = config

        # Copy tuning parameters to C++ extension

        self.config.set_tuning_params()

        # Read tensor list from file(s)

        if isinstance(self.config.model_path, str): model_path = [self.config.model_path]
        else: model_path = self.config.model_path

        # Read tensor list from file(s), and measure layer sizes

        load_keys = {}

        decoder_size = 0
        norm_size = 0
        head_size = 0

        for path in model_path:
            with safe_open(path, framework = "pt", device = "cpu") as f:
                for key in f.keys():

                    if _skip_key(key): continue

                    load_keys[key] = path

                    if key.startswith("model.layers.0."):
                        tensor_slice = f.get_slice(key)
                        shape = tensor_slice.get_shape()
                        decoder_size += math.prod(shape) * _layer_dtype_size(key)
                        del tensor_slice

                    if key.startswith("model.norm."):
                        tensor_slice = f.get_slice(key)
                        shape = tensor_slice.get_shape()
                        norm_size += math.prod(shape) * _layer_dtype_size(key)
                        del tensor_slice

                    if key.startswith("lm_head."):
                        tensor_slice = f.get_slice(key)
                        shape = tensor_slice.get_shape()
                        head_size += math.prod(shape) * _layer_dtype_size(key)
                        del tensor_slice

        # Begin auto mapping if enabled

        if self.config.auto_map is not None:

            self.config.device_map.embed_tokens = "cpu"
            self.config.device_map.layers = ["cuda:0"] + ["?"] * (self.config.num_hidden_layers - 1)

            # Assign layers automatically

            device_usage = 0
            device_index = 0
            layer_index_device = 0
            max_usage = self.config.auto_map[device_index] * (1024 ** 3)

            for layer in range(self.config.num_hidden_layers + 2):

                this_layer_size = decoder_size
                if layer == self.config.num_hidden_layers + 0: this_layer_size = norm_size
                elif layer == self.config.num_hidden_layers + 1: this_layer_size = head_size

                while device_usage + this_layer_size > max_usage:
                    device_index += 1
                    device_usage = 0
                    layer_index_device = 0
                    max_usage = self.config.auto_map[device_index] * (1024 ** 3)
                    if device_index >= len(self.config.auto_map): raise ValueError("Model too large for device allocation scheme.")

                target = f"cuda:{device_index}"
                if layer == self.config.num_hidden_layers + 0: self.config.device_map.norm = target
                elif layer == self.config.num_hidden_layers + 1: self.config.device_map.lm_head = target
                else: self.config.device_map.layers[layer] = f"cuda:{device_index}"

                device_usage += this_layer_size
                layer_index_device += 1

         # Load up to 1 GB of tensors at a time, closing and reopening the file in between each chunk

        max_dq_buffer_size = 0
        tensors = {}

        st_mem = 0
        MAX_ST_MEM = 1024**3
        f = None
        prev_path = ""
        for key, path in load_keys.items():

            device = self.config.device_map.map(key)

            if f is None or st_mem > MAX_ST_MEM or path != prev_path:
                if f is not None: del f
                f = safe_open(path, framework = "pt", device = "cpu")
                prev_path = path
                st_mem = 0

            tensor = f.get_tensor(key)
            size = tensor.numel() * tensor.element_size()
            st_mem += size

            if key.endswith(".scales"): tensor = tensor.half()
            if key == "lm_head.weight": tensor = tensor.float() if device == "cpu" else tensor.half()
            if key == "model.norm.weight": tensor = tensor.half()
            if key.endswith(".embed_tokens.weight"): tensor = tensor.half()
            if key.endswith(".input_layernorm.weight"): tensor = tensor.half()
            if key.endswith(".post_attention_layernorm.weight"): tensor = tensor.half()

            if device == "cpu": keep_tensor = tensor.clone()
            else: keep_tensor = tensor.to(device)
            del tensor

            if key.endswith(".qweight"): max_dq_buffer_size = max(max_dq_buffer_size, keep_tensor.numel() * 8)

            tensors[key] = keep_tensor

        del f

        # Head

        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias = False, device = "meta")
        self.lm_head.weight = nn.Parameter(tensors["lm_head.weight"])
        # self.lm_head_data = tensors["lm_head.weight"].transpose(0, 1).contiguous()

        # Token embeddings

        self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.hidden_size, self.config.pad_token_id, device = "meta")
        self.embed_tokens.weight = nn.Parameter(tensors["model.embed_tokens.weight"])
        with torch.no_grad():
            self.embed_tokens.weight[self.config.pad_token_id] = 0

        # Norm

        self.norm = ExLlamaRMSNorm(self.config, tensors, "model.norm.weight")

        # Prepare position embeddings for max seq length

        devs = self.config.device_map.get_layers_devs()

        self.sincos = {}
        for device in devs:

            inv_freq = 1.0 / (self.config.rotary_embedding_base ** (torch.arange(0, self.config.head_dim, 2, device = device).float() / self.config.head_dim))
            t = torch.arange(self.config.max_seq_len, device = device, dtype = torch.float32)
            if self.config.compress_pos_emb != 1.0: t /= self.config.compress_pos_emb

            freqs = torch.einsum("i,j->ij", t, inv_freq)
            emb = torch.cat((freqs, freqs), dim = -1)

            sin = emb.sin()[None, None, :, :].half()
            cos = emb.cos()[None, None, :, :].half()

            self.sincos[device] = (sin, cos)

        # Decoder layers

        modules = []
        device_layer_index = [0] * len(devs)

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

            device = self.config.device_map.layers[i]
            sin, cos = self.sincos[device]

            layer = ExLlamaDecoderLayer(self.config, tensors, f"model.layers.{i}", i, sin, cos)

            modules.append(layer)

        self.layers = modules

        # Prepare CUDA buffers

        self.buffers = []
        for dev in self.config.device_map.get_layers_devs():

            device_buffers = {}
            self.buffers.append(device_buffers)

            temp_state = torch.zeros((config.max_input_len, config.intermediate_size), dtype = torch.float16, device = dev)
            temp_mlp = torch.zeros((config.fused_mlp_thd * 2, config.intermediate_size), dtype = torch.float16, device = dev)
            temp_zeros_float = torch.zeros((1, 65536), dtype = torch.float32, device = dev)
            temp_dq = torch.zeros((1, max_dq_buffer_size), dtype = torch.float16, device = dev)

            device_buffers["temp_state"] = temp_state
            device_buffers["temp_mlp"] = temp_mlp
            device_buffers["temp_zeros_float"] = temp_zeros_float
            device_buffers["temp_dq"] = temp_dq

            cuda_ext.exllama_ext.prepare_buffers(torch.device(dev),
                                                 temp_state,
                                                 temp_mlp,
                                                 temp_zeros_float,
                                                 temp_dq)

        # Clear the cache

        torch.cuda.empty_cache()


    def forward(self,
                input_ids,
                cache,
                last_id_only = True,
                preprocess_only = False,
                lora = None,
                output_device = None,
                input_mask = None):

        q_len = input_ids.shape[-1]
        remaining_q_len = q_len
        bsz = input_ids.shape[0]

        assert input_mask is None or (input_mask.shape[-1] >= input_ids.shape[-1] and input_mask.shape[-2] == input_ids.shape[-2])

        # The buffers can only fit max_input_len tokens, so with larger batch sizes we reduce our work size correspondingly.

        effective_max_input_len = self.config.max_input_len // bsz

        # Split sequence

        result = None

        chunk_begin = 0
        while chunk_begin < q_len:

            # Limit chunk_size to max_input_len

            chunk_size = min(remaining_q_len, effective_max_input_len)

            # Limit chunk_size to keep size of attention operation <= max_attention_size, unless using flash-attn

            if not self.config.use_flash_attn_2 or chunk_begin > 0:

                past_len = cache.current_seq_len
                attn_size = (past_len + remaining_q_len) * remaining_q_len
                max_a = self.config.max_attention_size
                if attn_size > max_a:
                    cs = (math.sqrt(past_len ** 2 + 4 * max_a) - past_len) / 2
                    chunk_size = min(chunk_size, math.floor(cs))

            # Process chunk

            chunk_end = min(chunk_begin + chunk_size, q_len)

            _last_id_only = last_id_only
            _preprocess_only = preprocess_only or (chunk_end < q_len and last_id_only)

            r = self._forward(input_ids[:, chunk_begin : chunk_end],
                             cache,
                             _last_id_only,
                             _preprocess_only,
                             lora,
                             output_device,
                             input_mask)

            if not _preprocess_only:
                result = r if result is None else torch.cat((result, r), dim = 1)

            chunk_begin = chunk_end
            remaining_q_len -= chunk_size

        return result


    def _forward(self,
                 input_ids,
                 cache,
                 last_id_only = True,
                 preprocess_only = False,
                 lora = None,
                 output_device = None,
                 input_mask = None):

        # if torch.is_grad_enabled():
        #     raise ValueError("Forward pass called with gradients enabled. Back propagation is not supported yet.")
        with torch.no_grad():

            batch_size, seq_len = input_ids.shape
            past_len = cache.current_seq_len
            if output_device is None: output_device = input_ids.device

            buffer = ExLlamaBuffer(self.config)

            # Build attention mask on first device, copy to others if necessary

            devs = self.config.device_map.get_layers_devs()

            # if not self.config.use_flash_attn_2:

            if seq_len > 1 or input_mask is not None:

                attn_mask = torch.zeros(batch_size, 1, seq_len, past_len + seq_len, dtype = torch.float16, device = devs[0])
                attn_mask_triu = torch.triu(torch.full((seq_len - 1, seq_len - 1), -65504.))
                attn_mask[:, :, : seq_len - 1, past_len + 1: past_len + seq_len] = attn_mask_triu

                if input_mask is not None:

                    input_mask = input_mask[:, :past_len + seq_len]
                    input_mask = _move_tensor(input_mask, devs[0], "input_mask", self.config)
                    input_mask = torch.where(input_mask, 0, -65504.).half()
                    input_mask = input_mask.unsqueeze(1).unsqueeze(2)
                    attn_mask = torch.minimum(attn_mask, input_mask)

            else:

                attn_mask = None
                # attn_mask = torch.zeros(batch_size, 1, seq_len, seq_len + past_len, dtype = torch.float16, device = devs[0])

            buffer.attn_mask = attn_mask

            # else:
            #
            #     buffer.attn_mask = None

            # Embeddings
            # TODO: Allow passing input embeddings instead of IDs

            input_ids = _move_tensor(input_ids, self.config.device_map.embed_tokens, "input_ids", self.config)
            hidden_states = self.embed_tokens(input_ids)

            # Split buffers to devices

            buffers = {devs[0]: buffer}
            for device in devs[1:]:
                buffers[device] = buffer.to(device)

            # Decoder layers

            for i, decoder_layer in enumerate(self.layers):

                device = self.config.device_map.layers[i]
                hidden_states = _move_tensor(hidden_states, device, "hidden_states", self.config)

                hidden_states = decoder_layer.forward(hidden_states, cache, buffers[device], lora)

            cache.current_seq_len += seq_len

            # Early exit when we don't need logits

            if preprocess_only: return None

            # Norm

            hidden_states = _move_tensor(hidden_states, self.config.device_map.norm, "hidden_states", self.config)
            hidden_states = self.norm.forward(hidden_states, buffer)

            # Head

            if last_id_only: hidden_states = hidden_states[:, -1:, :].contiguous()
            if self.config.device_map.lm_head == "cpu": hidden_states = hidden_states.float()

            hidden_states = _move_tensor(hidden_states, self.config.device_map.lm_head, "hidden_states", self.config)
            logits = self.lm_head(hidden_states)
            # logits = cuda_ext.matmul_half(hidden_states, self.lm_head_data, cublas = False)

            logits = logits.float()
            logits = _move_tensor(logits, output_device, "logits", self.config)
            return logits


    # Free unmanaged resources allocated by the C++ extension. Call this before dereferencing the ExLlama object,
    # e.g. if you intend to create a new instance to load another model, but don't call it in a destructor that wraps
    # the object, since it relies on CUDA function calls and the CUDA context is one of the first things to go when
    # a PyTorch application terminates, before other managed objects are destroyed.

    def free_unmanaged(self):

        cuda_ext.exllama_ext.cleanup()