Support kv cache u8 precision

Author: riverlijunjie

src/common/transformations/src/transformations/common_optimizations/convert_pagedattn_inputs.cpp

@@ -107,7 +107,7 @@ ov::pass::ConvertPagedAttnInputs::ConvertPagedAttnInputs(const KVCacheConfig& co
         value_cache->set_element_type(value_cache_precision);
         bool status = false;
         if (pa_op->get_rt_info().count("num_k_heads") && pa_op->get_rt_info().count("k_head_size") &&
-            pa_op->get_rt_info().count("num_v_heads") && pa_op->get_rt_info().count("num_v_heads")) {
+            pa_op->get_rt_info().count("num_v_heads") && pa_op->get_rt_info().count("v_head_size")) {
             const auto key_cache_shape = init_cache_shape(pa_op->get_rt_info()["num_k_heads"].as<size_t>(),
                                                           pa_op->get_rt_info()["k_head_size"].as<size_t>(),
                                                           m_config.keyCacheBlockSize,

src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp

@@ -32,8 +32,13 @@ extern "C" _GENX_MAIN_ void KERNEL_NAME(
     const int32_t* block_indices [[type("svmptr_t")]],
     const int32_t* block_indices_begins [[type("svmptr_t")]],
     const int32_t* subsequence_begins [[type("svmptr_t")]],
+#if KV_CACHE_COMPRESSION_PER_TOKEN
+    uint8_t* key_cache [[type("svmptr_t")]],
+    uint8_t* value_cache [[type("svmptr_t")]],
+#else
     half* key_cache [[type("svmptr_t")]],
-    half* value_cache [[type("svmptr_t")]],    
+    half* value_cache [[type("svmptr_t")]],
+#endif   
     uint32_t key_pitch,
     uint32_t key_offset,
     uint32_t value_pitch,
@@ -84,14 +89,43 @@ extern "C" _GENX_MAIN_ void KERNEL_NAME(
 
     const uint block_offset = block_indices_begins[subsequence_idx] + current_block_idx;
 
+    #if KV_CACHE_COMPRESSION_PER_TOKEN
+    // Assume: K_HEAD_SIZE == K_HEAD_SIZE
+    auto quantize_and_store = [&](vector<half, K_HEAD_SIZE> data, uchar* out, uint out_offset, uint token_pos) {
+            uint scale_offset = out_offset + K_HEAD_SIZE * PAGED_ATTENTION_BLOCK_SIZE + token_pos * sizeof(half);
+            half max_val = cm_reduced_max<half>(data);
+            half min_val = cm_reduced_min<half>(data);
+            half scale_val = half(0.0);
+            half zp_val = half(0.0);
+            if(max_val == min_val) {
+                scale_val = half(0.0);
+                zp_val = max_val;
+            } else {
+                scale_val = 255.0 / (max_val - min_val);
+                zp_val = (0.0 - min_val) * scale_val;
+            }
+            vector<half, K_HEAD_SIZE>  dequant_data = cm_mul<half>(data, scale_val) + zp_val;
+            vector<uchar, K_HEAD_SIZE> data_u8 = cm_rnde<uchar, K_HEAD_SIZE>(dequant_data);
+            cm_ptr_store<uint32_t, K_HEAD_SIZE / 4>((uint32_t*)(out + out_offset + token_pos * K_HEAD_SIZE), 0, data_u8.format<uint32_t>());
+            half *out_scale_zp = (half*)(out + scale_offset);
+            out_scale_zp[0] = (max_val - min_val) / 255.0;
+            out_scale_zp[PAGED_ATTENTION_BLOCK_SIZE] = zp_val;
+    };
+    #endif
+
     {
         uint block_k_base_offset = (block_indices[block_offset] * KV_HEADS_NUM + head_idx) * ADJUSTED_K_HEAD_SIZE * PAGED_ATTENTION_BLOCK_SIZE;
         uint key_out_offset = block_k_base_offset + token_start_pos * K_HEAD_SIZE;
         uint key_in_offset = token_idx * key_pitch + head_idx * K_HEAD_SIZE + key_offset;
 
         vector<half, K_HEAD_SIZE> key_data;
         key_data.format<int>() = cm_ptr_load<int, K_HEAD_SIZE / 2>((int*)key, key_in_offset * (int)sizeof(half));
+
+        #if KV_CACHE_COMPRESSION_PER_TOKEN
+        quantize_and_store(key_data, (uchar*)key_cache, block_k_base_offset, token_start_pos);
+        #else
         cm_ptr_store<int, K_HEAD_SIZE / 2>((int*)key_cache, key_out_offset * (int)sizeof(half), key_data.format<int>());
+        #endif
     }
     {
         uint block_v_base_offset = (block_indices[block_offset] * KV_HEADS_NUM + head_idx) * ADJUSTED_V_HEAD_SIZE * PAGED_ATTENTION_BLOCK_SIZE;
@@ -106,6 +140,10 @@ extern "C" _GENX_MAIN_ void KERNEL_NAME(
 
         vector<half, V_HEAD_SIZE> value_data;
         value_data.format<int>() = cm_ptr_load<int, V_HEAD_SIZE / 2>((int*)value, value_in_offset * (int)sizeof(half));
+        #if KV_CACHE_COMPRESSION_PER_TOKEN
+        quantize_and_store(value_data, (uchar*)value_cache, block_v_base_offset, token_start_pos);
+        #else
         cm_ptr_store<int, V_HEAD_SIZE / 2>((int*)value_cache, value_out_offset * (int)sizeof(half), value_data.format<int>());
+        #endif
     }
 }

src/plugins/intel_gpu/src/graph/impls/cm/pa_kv_cache_update_ref.cm

@@ -26,15 +26,22 @@ namespace KERNEL_NAME {
 
 //extern "C" _GENX_MAIN_ void pa_multi_token(
 extern "C" _GENX_MAIN_ void KERNEL_NAME(
+    //query [q_len, num_heads, S]
     half* query [[type("svmptr_t")]],
-    half* key [[type("svmptr_t")]],
-    half* value [[type("svmptr_t")]],
+#if CMPA_KVCACHE_U8
+    int8_t* k_cache [[type("svmptr_t")]],
+    int8_t* v_cache [[type("svmptr_t")]],
+#else
+    half* k_cache [[type("svmptr_t")]],
+    half* v_cache [[type("svmptr_t")]],
+#endif
     int32_t* past_lens [[type("svmptr_t")]],
     int32_t* block_indices [[type("svmptr_t")]],
     int32_t* block_indices_begins [[type("svmptr_t")]],
     int32_t* subsequence_begins [[type("svmptr_t")]],
 #if SPARSE_BLOCK_SIZE > 1
     bool* sparse_block_mask [[type("svmptr_t")]],
+    bool* sparse_block_mask_wg [[type("svmptr_t")]],
 #endif
     half* output [[type("svmptr_t")]],
     int q_len) {
@@ -44,16 +51,26 @@ extern "C" _GENX_MAIN_ void KERNEL_NAME(
     constexpr int num_kv_heads = CMFLA_NUM_KV_HEADS;
     constexpr int pa_block_sz = CMPA_BLOCK_SZ;
     //# query [q_len, num_heads, S]
-    //#   key [kv_len, num_heads, S]
-    //# value [kv_len, num_heads, S]
-    //# sparse_block_mask [num_heads, q_blocks, kv_blocks]
+    //# k_cache [kv_len, num_heads, S]
+    //# v_cache [kv_len, num_heads, S]
+#if CMPA_KVCACHE_U8
+    constexpr uint K_SLM_SIZE = (4*kv_step * head_size * sizeof(half));
+    constexpr uint V_SLM_SIZE = (4*kv_step * head_size * sizeof(half));
+    constexpr uint Q_SLM_SIZE = 0;//(q_step * head_size * sizeof(half)) * local_size;
+
+    cm_slm_init(K_SLM_SIZE + V_SLM_SIZE + Q_SLM_SIZE);
 
+    auto slm_K = cm_slm_alloc(K_SLM_SIZE);
+    auto slm_V = cm_slm_alloc(V_SLM_SIZE);
+
+#endif
     auto batch = cm_group_id(0);
     auto h = cm_group_id(1);
     auto hkv = h / (num_heads/num_kv_heads);
     auto wg_id = cm_group_id(2); // each work-group handles a sequence
     auto wg_local_id = cm_local_id(2);
     int local_size = cm_local_size(2);
+
     int q_start_sg, kv_start, kv_seq_len, q_len_sg;
 
     // multiple work-groups are required to split a sequence,
@@ -91,57 +108,71 @@ extern "C" _GENX_MAIN_ void KERNEL_NAME(
         ---------------------------------
         each grid can be [q_len_per_trunk, q_len_per_trunk].
         For each trunk, [q_len_per_trunk, past_q_lens] must be calculated. Such as: `20`,`21`. but for the 22,
-        causal mask optimization can be applied. different wgs would has different kv stop.
+        casual mask optimization can be applied. differnt wgs would has different kv stop.
         //todo:kv_stop is wg level, should we change to sg level?
+               sglevel would cause sgs in one wg diverge. so leave for now. also one wg has same kvstop makes eaiser for kv copying/loading into SLM/cache.
     */
         kv_stop = (wg_id + 1) * wg_seq_len + past_q_lens;
         if (kv_stop > kv_seq_len) kv_stop = kv_seq_len;
     }
-
-    // printf("wg:%d.%d  q: %d, +%d   kv: %d, +%d, %d\n", wg_id, wg_local_id, q_start_sg, q_len_sg, kv_start, kv_seq_len, kv_stop);
-    // qkv fused
-    // constexpr uint num_total_heads = num_heads + num_kv_heads * 2;
-    // uint q_offset = (q_start*num_total_heads + h)*head_size;
-    // uint k_offset = (kv_start*num_total_heads + num_heads + hkv)*head_size;
-    // uint v_offset = (kv_start*num_total_heads + num_heads + num_kv_heads + hkv)*head_size;
+    // printf("###########wg:%d.%d  q: %d, +%d   kv: %d, +%d, kvstop:%d\n", wg_id, wg_local_id, q_start_sg, q_len_sg, kv_start, kv_seq_len, kv_stop);
 
     //Q/O[B, L, H, S]
     uint q_offset = (q_start_sg*num_heads + h)*head_size;
-    uint o_offset = (q_start_sg*num_heads + h)*head_size;
-
-    //K/V[block_num, kv_heads, block_sz, head_sz]
-    uint k_offset = hkv*head_size*pa_block_sz;
-    uint v_offset = hkv*head_size*pa_block_sz;
 
 #if SPARSE_BLOCK_SIZE > 1
     //# sparse_block_mask [num_heads, q_blocks, kv_blocks]
     auto q_start_block = q_start_sg/ SPARSE_BLOCK_SIZE;
     int q_blocks = (q_len + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE;
     int kv_blocks = (kv_seq_len + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE;
+    //[self.num_heads, q_block_num, kv_block_num]
     bool* block_mask_base = sparse_block_mask + (h * q_blocks + q_start_block)*kv_blocks;
+    //[self.num_heads, wg_count_along_query, kv_block_num)]
+    bool* wg_block_mask_base = sparse_block_mask_wg + (h * cm_group_count(2) + wg_id)*kv_blocks;
     // printf("wg:%d.%d  q: %d, +%d   kv: %d, +%d, %d, x-attn: %d, %dx%d, %p, %p\n", wg_id, wg_local_id, q_start_sg, q_len_sg, kv_start, kv_seq_len, kv_stop, q_start_block, q_blocks, kv_blocks, sparse_block_mask, block_mask_base);
 #endif
 
-#if USE_LSC == 1
-    pa_kernel_lsc_prefetch<is_causal, num_heads, num_kv_heads, head_size, 0, 16>(
-                                wg_local_id,
-                                q_start_sg, //q_start for SG,
-                                kv_stop,
-                                q_len_sg, //q_step,
-                                kv_seq_len, //kv_len, not used for now
-                                reinterpret_cast<svmptr_t>(query + q_offset),
-                                reinterpret_cast<svmptr_t>(key + k_offset),
-                                reinterpret_cast<svmptr_t>(value + v_offset),
+#if CMPA_KVCACHE_U8
+    uint kv_offset = hkv*(head_size+4)*pa_block_sz;
+    pa_lsc_u8<is_causal, num_heads, num_kv_heads, head_size, 0>(
+                            slm_K,
+                            slm_V,
+                            wg_local_id,
+                            local_size,
+                            q_start_sg, //q_start for SG,
+                            kv_stop,
+                            q_len_sg, //q_step,
+                            kv_seq_len, //kv_len,
+                            reinterpret_cast<svmptr_t>(query + q_offset),
+                            reinterpret_cast<svmptr_t>(k_cache + kv_offset),
+                            reinterpret_cast<svmptr_t>(v_cache + kv_offset),
 #if SPARSE_BLOCK_SIZE > 1
-                                reinterpret_cast<svmptr_t>(block_mask_base),
+                            reinterpret_cast<svmptr_t>(block_mask_base),
+                            reinterpret_cast<svmptr_t>(wg_block_mask_base),
+
 #endif
-                                reinterpret_cast<svmptr_t>(output + o_offset),
-                                past_q_lens,
-                                block_indices);
+                            reinterpret_cast<svmptr_t>(output + q_offset),
+                            past_q_lens,
+                            block_indices);
 #else
-    static_assert(0);
+    uint kv_offset = hkv*head_size*pa_block_sz;
+    pa_kernel_lsc_prefetch_f16<is_causal, num_heads, num_kv_heads, head_size, 0, 16>(
+                            wg_local_id,
+                            q_start_sg, //q_start for SG,
+                            kv_stop,
+                            q_len_sg, //q_step,
+                            kv_seq_len, //kv_len,
+                            reinterpret_cast<svmptr_t>(query + q_offset),
+                            reinterpret_cast<svmptr_t>(k_cache + kv_offset),
+                            reinterpret_cast<svmptr_t>(v_cache + kv_offset),
+#if SPARSE_BLOCK_SIZE > 1
+                            reinterpret_cast<svmptr_t>(block_mask_base),
+                            reinterpret_cast<svmptr_t>(wg_block_mask_base),
 
+#endif
+                            reinterpret_cast<svmptr_t>(output + q_offset),
+                            past_q_lens,
+                            block_indices);
 #endif
 }
-
-}  // NAMESPACE
+} // namespace KERNEL_NAME

src/plugins/intel_gpu/src/graph/impls/cm/pa_multi_token.cm

@@ -45,9 +45,6 @@
 
 #define KV_PARTITION_STEP_NUM  (KV_PARTITION_SIZE / KV_STEP)
 
-#define KV_SCALE_ZP_SIZE 0 // 4: scale/zp size
-
-
 #define DEBUG_ENABLE 0
 #if DEBUG_ENABLE
 template<typename T, int M, int N>
@@ -103,7 +100,7 @@ void show(vector<T, N> vec) {
 
 //prepack [K, N] to [K/2, N, 2] layout.
 template <typename T1, typename T2, int K, int N>
-inline void prepackAsVNNIWidth2(matrix_ref<T1, K, N> input, matrix_ref<T2, K/2, N*2> out) {
+inline void prepack_to_VNNI_W2(matrix_ref<T1, K, N> input, matrix_ref<T2, K/2, N*2> out) {
     #pragma unroll
     for (int r = 0; r < K/2; r++) {
         out.row(r).select<N, 2>(0) = input.row(r*2);
@@ -498,7 +495,7 @@ extern "C" _GENX_MAIN_ void KERNEL_NAME(
                             VmatNormal[r] = 0;
                         }
                     }
-                    prepackAsVNNIWidth2(VmatNormal, Vmat.format<half, REG_K/2, REG_N*2>());
+                    prepack_to_VNNI_W2(VmatNormal, Vmat.format<half, REG_K/2, REG_N*2>());
                 #else
                     cm_load<lsc::VNNI>(Vmat[0].format<half>(), b2dV.set_block_y(kv_pos));
                 #endif

src/plugins/intel_gpu/src/graph/impls/cm/pa_single_token.cm

@@ -27,7 +27,7 @@ struct PagedAttentionImplementationManager : public ImplementationManager {
         };
         static constexpr std::array supported_kv_types = {
             ov::element::f16,
-            // ov::element::i8,
+            ov::element::i8,
         };
 
         auto& engine = node.get_program().get_engine();

src/plugins/intel_gpu/src/graph/impls/cm/paged_attention.hpp

@@ -66,10 +66,25 @@ inline size_t get_kv_len(const RuntimeParams& params, const PagedAttentionStage&
     return 0;  // Fallback case, should not be reached
 }
 
+inline size_t get_input_kv_len(const RuntimeParams& params) {
+    auto key_shape = params.input_layouts[PagedAttentionInputIdx::KEY].get_shape();
+    const size_t kv_len = key_shape[key_shape.size() - 2];
+    return kv_len;
+}
+
 inline size_t get_aligned_kv_len(const size_t kv_len) {
     return (kv_len + PA_KV_CACHE_BLOCK_SIZE - 1) / PA_KV_CACHE_BLOCK_SIZE * PA_KV_CACHE_BLOCK_SIZE;
 }
 
+inline bool get_kv_compressed(const RuntimeParams& params) {
+    auto key_cache_layout = params.input_layouts[PagedAttentionInputIdx::KEY_CACHE];
+    if (data_type_traits::is_i8_u8(key_cache_layout.data_type)) {
+        return true;
+    } else {
+        return false;
+    }
+}
+
 int64_t get_aligned_seq_len(const kernel_impl_params& impl_param, const PagedAttentionStage& stage, int64_t target_seq_len_block_size = 16) {
     // Since at prefill stage Q, K, V inputs may contain multiple sequences with arbitrary
     // target sequence lengths each (shape is [sequences_num * target_seq_len, num_heads * head_size]),
@@ -268,10 +283,18 @@ JitConstants PagedAttentionGeneratorKVCacheUpdate::get_jit_constants(const kerne
     jit.make("KV_HEADS_NUM", desc->kv_heads_num);
     jit.make("K_HEAD_SIZE", desc->k_head_size);
     jit.make("V_HEAD_SIZE", desc->v_head_size);
-    jit.make("ADJUSTED_K_HEAD_SIZE", desc->k_head_size);
-    jit.make("ADJUSTED_V_HEAD_SIZE", desc->v_head_size);
     jit.make("PAGED_ATTENTION_BLOCK_SIZE", PA_KV_CACHE_BLOCK_SIZE);
 
+    if (get_kv_compressed(params)) {
+        jit.make("KV_CACHE_COMPRESSION_PER_TOKEN", 1);
+        jit.make("ADJUSTED_K_HEAD_SIZE", desc->k_head_size + 4);
+        jit.make("ADJUSTED_V_HEAD_SIZE", desc->v_head_size + 4);
+    } else {
+        jit.make("KV_CACHE_COMPRESSION_PER_TOKEN", 0);
+        jit.make("ADJUSTED_K_HEAD_SIZE", desc->k_head_size);
+        jit.make("ADJUSTED_V_HEAD_SIZE", desc->v_head_size);
+    }
+
     return jit;
 }
 
@@ -302,7 +325,8 @@ DispatchDataFunc PagedAttentionGeneratorKVCacheUpdate::get_dispatch_data_func()
         const auto desc = params.typed_desc<paged_attention>();
         // auto rtp = static_cast<PagedAttentionRuntimeParams*>(rt_params);
 
-        const size_t kv_len = get_max_context_len(params);
+        // const size_t kv_len = get_max_context_len(params);
+        const size_t kv_len = get_input_kv_len(params);
         const size_t kv_heads_num = desc->kv_heads_num;
         const size_t wg_count = (kv_len + WG_SIZE - 1) / WG_SIZE;
 
@@ -372,7 +396,8 @@ DispatchDataFunc PagedAttentionGeneratorKVCacheUpdate::get_dispatch_data_func()
         if (DEBUG_ENABLED) {  // Debug
             std::cout << "PagedAttentionGeneratorKVCacheUpdate::get_dispatch_data_func: "
                       << "kv_len: " << kv_len << ", key_pitch: " << key_pitch << ", key_offset: " << key_offset << ", value_pitch: " << value_pitch
-                      << ", value_offset: " << value_offset << ", "<< std::endl;
+                      << ", value_offset: " << value_offset << ", gws: [" << wgs.global[0] << ", " << wgs.global[1] << ", " << wgs.global[2] << "]"
+                      << ", lws: [" << wgs.local[0] << ", " << wgs.local[1] << ", " << wgs.local[2] << "]" << std::endl;
         }
 
         // TODO: support multiple sequences
@@ -429,6 +454,12 @@ JitConstants PagedAttentionGeneratorMultiToken::get_jit_constants(const kernel_i
     jit.make("CMPA_BLOCK_SZ", PA_KV_CACHE_BLOCK_SIZE);
     jit.make("SPARSE_BLOCK_SIZE", xattn_block_size);
     jit.make("Q_STEP", get_q_step(xe_arch, true));
+
+    if (get_kv_compressed(params)) {
+        jit.make("CMPA_KVCACHE_U8", 1);
+    } else {
+        jit.make("CMPA_KVCACHE_U8", 0);
+    }
     // for (auto& it : jit) {
     //     std::cout << "\tjit[" << it.name << "] = " << it.value << std::endl;
     // }
@@ -509,7 +540,11 @@ JitConstants PagedAttentionGeneratorSingleToken::get_jit_constants(const kernel_
     jit.make("KV_HEADS_NUM", desc->kv_heads_num);
     jit.make("Q_STEP", get_q_step(xe_arch, true));
 
-    jit.make("KV_CACHE_COMPRESSION", 0);
+    if (get_kv_compressed(params)) {
+        jit.make("KV_CACHE_COMPRESSION", 1);
+    } else {
+        jit.make("KV_CACHE_COMPRESSION", 0);
+    }
 
     return jit;
 }

src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp

@@ -26,7 +26,7 @@ struct PagedAttentionOpt : public ImplementationManager {
         };
         static constexpr std::array supported_kv_types = {
             #if ENABLE_PA_CM_PATH
-            ov::element::i8,
+            ov::element::f32,
             #else
             ov::element::f32,
             ov::element::f16,