Add Q6 cache mode

Author: turboderp

examples/chat.py

@@ -8,6 +8,7 @@
     ExLlamaV2Cache,
     ExLlamaV2Cache_8bit,
     ExLlamaV2Cache_Q4,
+    ExLlamaV2Cache_Q6,
     ExLlamaV2Cache_Q8,
     ExLlamaV2Tokenizer,
     model_init,
@@ -55,6 +56,7 @@
 
 parser.add_argument("-c8", "--cache_8bit", action = "store_true", help = "Use 8-bit (FP8) cache")
 parser.add_argument("-cq4", "--cache_q4", action = "store_true", help = "Use Q4 cache")
+parser.add_argument("-cq6", "--cache_q6", action = "store_true", help = "Use Q6 cache")
 parser.add_argument("-cq8", "--cache_q8", action = "store_true", help = "Use Q8 cache")
 
 parser.add_argument("-ngram", "--ngram_decoding", action = "store_true", help = "Use n-gram speculative decoding")
@@ -130,6 +132,8 @@
         draft_cache = ExLlamaV2Cache_8bit(draft_model)
     elif args.cache_q4:
         draft_cache = ExLlamaV2Cache_Q4(draft_model)
+    elif args.cache_q6:
+        draft_cache = ExLlamaV2Cache_Q6(draft_model)
     elif args.cache_q8:
         draft_cache = ExLlamaV2Cache_Q8(draft_model)
     else:
@@ -141,6 +145,8 @@
     cache = ExLlamaV2Cache_8bit(model, lazy = not model.loaded)
 elif args.cache_q4:
     cache = ExLlamaV2Cache_Q4(model, lazy = not model.loaded)
+elif args.cache_q6:
+    cache = ExLlamaV2Cache_Q6(model, lazy=not model.loaded)
 elif args.cache_q8:
     cache = ExLlamaV2Cache_Q8(model, lazy = not model.loaded)
 else:

exllamav2/__init__.py

@@ -4,6 +4,7 @@
 from exllamav2.cache import ExLlamaV2CacheBase
 from exllamav2.cache import ExLlamaV2Cache
 from exllamav2.cache import ExLlamaV2Cache_Q4
+from exllamav2.cache import ExLlamaV2Cache_Q6
 from exllamav2.cache import ExLlamaV2Cache_Q8
 from exllamav2.cache import ExLlamaV2Cache_8bit
 from exllamav2.config import ExLlamaV2Config

exllamav2/cache.py

@@ -25,7 +25,8 @@ class ExLlamaV2CacheBase:
     head_dim: int
 
     dtype: torch.dtype
-    weights_per_element: int
+    weights_per_element_k: int
+    weights_per_element_v: int
     has_scales: bool
 
 
@@ -34,14 +35,16 @@ def __init__(self,
                  batch_size: int,
                  max_seq_len: int,
                  dtype: torch.dtype,
-                 weights_per_element: int,
+                 weights_per_element_k: int,
+                 weights_per_element_v: 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.weights_per_element_k = weights_per_element_k
+        self.weights_per_element_v = weights_per_element_v
         self.has_scales = has_scales
 
         self.key_states = []
@@ -55,7 +58,8 @@ def __init__(self,
 
         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_wk = (self.batch_size, self.max_seq_len, self.num_key_value_heads, self.head_dim // self.weights_per_element_k)
+        self.shape_wv = (self.batch_size, self.max_seq_len, self.num_key_value_heads, self.head_dim // self.weights_per_element_v)
         self.shape_s = (self.batch_size, self.max_seq_len, self.num_key_value_heads, self.head_dim // 32)
 
 
@@ -74,8 +78,8 @@ def create_state_tensors(self,
 
                 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()
+                    p_key_states = torch.zeros(self.shape_wk, dtype = self.dtype, device = device).contiguous()
+                    p_value_states = torch.zeros(self.shape_wv, 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()
@@ -115,8 +119,8 @@ def update_cache_tensors(self):
                 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()
+            p_key_states = torch.zeros(self.shape_wk, dtype = self.dtype, device = v).contiguous()
+            p_value_states = torch.zeros(self.shape_wv, dtype = self.dtype, device = v).contiguous()
             self.key_states[k] = p_key_states
             self.value_states[k] = p_value_states
             if self.has_scales:
@@ -220,7 +224,7 @@ def __init__(self,
                  copy_from: ExLlamaV2Cache | None = None,
                  lazy: bool = False):
 
-        super().__init__(model, batch_size, max_seq_len, torch.half, 1, False)
+        super().__init__(model, batch_size, max_seq_len, torch.half, 1, 1, False)
 
         self.create_state_tensors(copy_from, lazy)
 
@@ -280,7 +284,7 @@ def __init__(self,
                  copy_from: ExLlamaV2Cache_8bit | None = None,
                  lazy: bool = False):
 
-        super().__init__(model, batch_size, max_seq_len, torch.uint8, 1, False)
+        super().__init__(model, batch_size, max_seq_len, torch.uint8, 1, 1, False)
 
         self.create_state_tensors(copy_from, lazy)
 
@@ -365,9 +369,10 @@ def __init__(self,
                  max_seq_len: int = -1,
                  copy_from: ExLlamaV2Cache_Q4 | None = None,
                  lazy: bool = False,
-                 weights_per_byte: int = -1):
+                 weights_per_byte_k: int = -1,
+                 weights_per_byte_v: int = -1):
 
-        super().__init__(model, batch_size, max_seq_len, torch.uint8, weights_per_byte, True)
+        super().__init__(model, batch_size, max_seq_len, torch.uint8, weights_per_byte_k, weights_per_byte_v, True)
         cfg = self.model.config
 
         self.create_state_tensors(copy_from, lazy)
@@ -607,18 +612,31 @@ def __init__(self,
                  copy_from: ExLlamaV2Cache_Q4 | None = None,
                  lazy: bool = False):
 
-        super().__init__(model, batch_size, max_seq_len, copy_from, lazy, 2)
+        super().__init__(model, batch_size, max_seq_len, copy_from, lazy, 2, 2)
         self.wbits = 4
 
 
+class ExLlamaV2Cache_Q6(ExLlamaV2Cache_Q):
+
+    def __init__(self,
+                 model: ExLlamaV2,
+                 batch_size: int = 1,
+                 max_seq_len: int = -1,
+                 copy_from: ExLlamaV2Cache_Q6 | None = None,
+                 lazy: bool = False):
+
+        super().__init__(model, batch_size, max_seq_len, copy_from, lazy, 1, 2)
+        self.wbits = 6
+
+
 class ExLlamaV2Cache_Q8(ExLlamaV2Cache_Q):
 
     def __init__(self,
                  model: ExLlamaV2,
                  batch_size: int = 1,
                  max_seq_len: int = -1,
-                 copy_from: ExLlamaV2Cache_Q4 | None = None,
+                 copy_from: ExLlamaV2Cache_Q8 | None = None,
                  lazy: bool = False):
 
-        super().__init__(model, batch_size, max_seq_len, copy_from, lazy, 1)
+        super().__init__(model, batch_size, max_seq_len, copy_from, lazy, 1, 1)
         self.wbits = 8

exllamav2/exllamav2_ext/cuda/cache.cu

@@ -119,7 +119,7 @@ void array_fp8_to_fp16_cuda(const unsigned char* pIn, half* pOut, int stride, in
 
 // -------------- FP16 -> Q
 
-template <int wbits>
+template <int wbits_k, int wbits_v>
 __global__ void fp16_to_q_kv_paged_kernel
 (
     const half* __restrict__ k_in,
@@ -172,11 +172,14 @@ __global__ void fp16_to_q_kv_paged_kernel
     {
         int j = i + blockIdx.y * BLOCKSIZE_Q;
         if (j >= block_b) continue;
-        fp16_to_q<wbits>(t, in, out, scales, j, cal, dim);
+        if (kv)
+            fp16_to_q<wbits_v>(t, in, out, scales, j, cal, dim);
+        else
+            fp16_to_q<wbits_k>(t, in, out, scales, j, cal, dim);
     }
 }
 
-template <int wbits>
+template <int wbits_k, int wbits_v>
 __global__ void fp16_to_q_kv_kernel
 (
     const half* __restrict__ k_in,
@@ -193,13 +196,17 @@ __global__ void fp16_to_q_kv_kernel
 )
 {
     int t = threadIdx.x;
-    const half* in = blockIdx.z ? v_in : k_in;
-    unsigned char* out = blockIdx.z ? v_out : k_out;
-    half* scales = blockIdx.z ? v_scales : k_scales;
-    const half* cal = blockIdx.z ? cal_v : cal_k;
+    int kv = blockIdx.z & 1;
+    const half* in = kv ? v_in : k_in;
+    unsigned char* out = kv ? v_out : k_out;
+    half* scales = kv ? v_scales : k_scales;
+    const half* cal = kv ? cal_v : cal_k;
     int block_offset = (offset + blockIdx.y * stride + blockIdx.x * BLOCKSIZE_Q);
 
-    fp16_to_q<wbits>(t, in, out, scales, block_offset, cal, dim);
+    if (kv)
+        fp16_to_q<wbits_v>(t, in, out, scales, block_offset, cal, dim);
+    else
+        fp16_to_q<wbits_k>(t, in, out, scales, block_offset, cal, dim);
 }
 
 void array_fp16_to_q_kv_paged_cuda
@@ -229,7 +236,17 @@ void array_fp16_to_q_kv_paged_cuda
     gridDim.z = batch_size * 2;
 
     if (wbits == 4)
-        fp16_to_q_kv_paged_kernel<4><<<gridDim, blockDim>>>
+        fp16_to_q_kv_paged_kernel<4, 4><<<gridDim, blockDim>>>
+        (
+            k_in, k_out, k_scales,
+            v_in, v_out, v_scales,
+            cache_seqlens, block_table,
+            pages_per_seq, page_size,
+            dim, q_len,
+            cal_k, cal_v
+        );
+    else if (wbits == 6)
+        fp16_to_q_kv_paged_kernel<8, 4><<<gridDim, blockDim>>>
         (
             k_in, k_out, k_scales,
             v_in, v_out, v_scales,
@@ -239,7 +256,7 @@ void array_fp16_to_q_kv_paged_cuda
             cal_k, cal_v
         );
     else if (wbits == 8)
-        fp16_to_q_kv_paged_kernel<8><<<gridDim, blockDim>>>
+        fp16_to_q_kv_paged_kernel<8, 8><<<gridDim, blockDim>>>
         (
             k_in, k_out, k_scales,
             v_in, v_out, v_scales,
@@ -275,14 +292,21 @@ void array_fp16_to_q_kv_cuda
     gridDim.z = v_in ? 2 : 1;
 
     if (wbits == 4)
-        fp16_to_q_kv_kernel<4><<<gridDim, blockDim>>>(
+        fp16_to_q_kv_kernel<4, 4><<<gridDim, blockDim>>>(
+            k_in, k_out, k_scales,
+            v_in, v_out, v_scales,
+            dim, offset, stride,
+            cal_k, cal_v
+        );
+    else if (wbits == 6)
+        fp16_to_q_kv_kernel<8, 4><<<gridDim, blockDim>>>(
             k_in, k_out, k_scales,
             v_in, v_out, v_scales,
             dim, offset, stride,
             cal_k, cal_v
         );
     else if (wbits == 8)
-        fp16_to_q_kv_kernel<8><<<gridDim, blockDim>>>(
+        fp16_to_q_kv_kernel<8, 8><<<gridDim, blockDim>>>(
             k_in, k_out, k_scales,
             v_in, v_out, v_scales,
             dim, offset, stride,
@@ -292,7 +316,7 @@ void array_fp16_to_q_kv_cuda
 
 // --------------- Q -> FP16
 
-template <int wbits>
+template <int wbits_k, int wbits_v>
 __global__ void q_to_fp16_kv_paged_kernel
 (
     const unsigned char* __restrict__ k_in,
@@ -342,11 +366,14 @@ __global__ void q_to_fp16_kv_paged_kernel
     {
         int j = i + blockIdx.y * BLOCKSIZE_Q;
         if (j >= block_b) continue;
-        q_to_fp16<wbits>(t, in, scales, out, j, cal, dim);
+        if (kv)
+            q_to_fp16<wbits_v>(t, in, scales, out, j, cal, dim);
+        else
+            q_to_fp16<wbits_k>(t, in, scales, out, j, cal, dim);
     }
 }
 
-template <int wbits>
+template <int wbits_k, int wbits_v>
 __global__ void q_to_fp16_kv_kernel
 (
     const unsigned char* __restrict__ k_in,
@@ -363,13 +390,17 @@ __global__ void q_to_fp16_kv_kernel
 )
 {
     int t = threadIdx.x;
-    const unsigned char* in = blockIdx.z ? v_in : k_in;
-    const half* scales = blockIdx.z ? v_scales : k_scales;
-    half* out = blockIdx.z ? v_out : k_out;
-    const half* cal = blockIdx.z ? cal_v : cal_k;
+    int kv = blockIdx.z & 1;
+    const unsigned char* in = kv ? v_in : k_in;
+    const half* scales = kv ? v_scales : k_scales;
+    half* out = kv ? v_out : k_out;
+    const half* cal = kv ? cal_v : cal_k;
     int block_offset = (offset + blockIdx.y * stride + blockIdx.x * BLOCKSIZE_Q);
 
-    q_to_fp16<wbits>(t, in, scales, out, block_offset, cal, dim);
+    if (kv)
+        q_to_fp16<wbits_v>(t, in, scales, out, block_offset, cal, dim);
+    else
+        q_to_fp16<wbits_k>(t, in, scales, out, block_offset, cal, dim);
 }
 
 void array_q_to_fp16_kv_paged_cuda
@@ -398,7 +429,17 @@ void array_q_to_fp16_kv_paged_cuda
     gridDim.z = batch_size * 2;
 
     if (wbits == 4)
-        q_to_fp16_kv_paged_kernel<4><<<gridDim, blockDim>>>
+        q_to_fp16_kv_paged_kernel<4, 4><<<gridDim, blockDim>>>
+        (
+            k_in, k_scales, k_out,
+            v_in, v_scales, v_out,
+            cache_seqlens, block_table,
+            pages_per_seq, page_size,
+            dim,
+            cal_k, cal_v
+        );
+    else if (wbits == 6)
+        q_to_fp16_kv_paged_kernel<8, 4><<<gridDim, blockDim>>>
         (
             k_in, k_scales, k_out,
             v_in, v_scales, v_out,
@@ -408,7 +449,7 @@ void array_q_to_fp16_kv_paged_cuda
             cal_k, cal_v
         );
     else if (wbits == 8)
-        q_to_fp16_kv_paged_kernel<8><<<gridDim, blockDim>>>
+        q_to_fp16_kv_paged_kernel<8, 8><<<gridDim, blockDim>>>
         (
             k_in, k_scales, k_out,
             v_in, v_scales, v_out,
@@ -444,14 +485,21 @@ void array_q_to_fp16_kv_cuda
     gridDim.z = v_in ? 2 : 1;
 
     if (wbits == 4)
-        q_to_fp16_kv_kernel<4><<<gridDim, blockDim>>>(
+        q_to_fp16_kv_kernel<4, 4><<<gridDim, blockDim>>>(
+            k_in, k_scales, k_out,
+            v_in, v_scales, v_out,
+            dim, offset, stride,
+            cal_k, cal_v
+        );
+    else if (wbits == 6)
+        q_to_fp16_kv_kernel<8, 4><<<gridDim, blockDim>>>(
             k_in, k_scales, k_out,
             v_in, v_scales, v_out,
             dim, offset, stride,
             cal_k, cal_v
         );
     else if (wbits == 8)
-        q_to_fp16_kv_kernel<8><<<gridDim, blockDim>>>(
+        q_to_fp16_kv_kernel<8, 8><<<gridDim, blockDim>>>(
             k_in, k_scales, k_out,
             v_in, v_scales, v_out,
             dim, offset, stride,

exllamav2/exllamav2_ext/ext_cache.cpp

@@ -107,7 +107,7 @@ void fp16_to_q_kv
     TORCH_CHECK_SHAPES(k_in, 0, v_in, 0, 1);
     TORCH_CHECK_SHAPES(k_in, 1, v_in, 1, 1);
     TORCH_CHECK_SHAPES(k_in, 2, v_in, 2, 1);
-    TORCH_CHECK_SHAPES(k_in, 3, v_in, 3, 1);
+//    TORCH_CHECK_SHAPES(k_in, 3, v_in, 3, 1);
 
     if (!cal_k.device().is_meta())
         TORCH_CHECK_SHAPES_OPT(cal_k, 0, k_in, 2, 1);
@@ -207,7 +207,7 @@ void q_to_fp16_kv
     TORCH_CHECK_SHAPES(k_in, 0, v_in, 0, 1);
     TORCH_CHECK_SHAPES(k_in, 1, v_in, 1, 1);
     TORCH_CHECK_SHAPES(k_in, 2, v_in, 2, 1);
-    TORCH_CHECK_SHAPES(k_in, 3, v_in, 3, 1);
+//    TORCH_CHECK_SHAPES(k_in, 3, v_in, 3, 1);
 
     if (!cal_k.device().is_meta())
         TORCH_CHECK_SHAPES_OPT(cal_k, 0, k_out, 2, 1);