[GPU] Add the capability for KV cache to update past KV

src/plugins/intel_gpu/include/intel_gpu/graph/kernel_impl_params.hpp

@@ -79,6 +79,10 @@ struct kernel_impl_params final {
     std::vector<size_t> output_size;
     std::vector<size_t> img_size;
 
+    // KV cache trim length - set at runtime during shape inference
+    // Marked as mutable to allow modification even when kernel_impl_params is passed as const reference
+    mutable int64_t kv_cache_trim_length = 0;
+
     std::map<size_t, size_t> in_port_to_shape_info_offset = {};
     std::map<size_t, size_t> out_port_to_shape_info_offset = {};
 

src/plugins/intel_gpu/include/intel_gpu/op/kv_cache.hpp

@@ -33,6 +33,16 @@ class KVCache : public ov::op::Op, public ov::op::util::VariableExtension {
             int64_t gather_axis,
             const ov::element::Type output_type = ov::element::dynamic);
 
+    /// KVcache with reorder for tree-based speculative decoding
+    KVCache(const Output<Node>& past,
+            const Output<Node>& new_token_data,
+            const Output<Node>& past_seq_len,
+            const Output<Node>& dst_idx,
+            const Output<Node>& update_data,
+            const std::shared_ptr<ov::op::util::Variable>& past_values,
+            int64_t concat_axis,
+            const ov::element::Type output_type = ov::element::dynamic);
+
     bool visit_attributes(ov::AttributeVisitor& visitor) override;
 
     void validate_and_infer_types() override;
@@ -51,6 +61,11 @@ class KVCache : public ov::op::Op, public ov::op::util::VariableExtension {
     void set_gather_axis(int64_t axis) { m_gather_axis = axis; }
 
     bool get_indirect() const { return m_indirect; }
+    bool get_update_kv() const { return m_update_kv; }
+    void set_update_kv(bool update_kv) { m_update_kv = update_kv; }
+
+    uint64_t get_trim_length() const { return m_trim_length; }
+    void set_trim_length(uint64_t trim_length) { m_trim_length = trim_length; }
 
 protected:
     KVCache(const OutputVector& inputs,
@@ -63,6 +78,8 @@ class KVCache : public ov::op::Op, public ov::op::util::VariableExtension {
     int64_t m_concat_axis = 0;
     int64_t m_gather_axis = 0;
     bool m_indirect = false;
+    bool m_update_kv = false;
+    uint64_t m_trim_length = 0;
 
     ov::element::Type m_output_type;
 };

src/plugins/intel_gpu/include/intel_gpu/primitives/kv_cache.hpp

@@ -27,18 +27,20 @@ struct kv_cache : public primitive_base<kv_cache> {
              const ov::op::util::VariableInfo& variable_info,
              const int64_t concat_axis,
              const int64_t gather_axis,
-             const bool indirect)
+             const bool indirect,
+             const bool update_kv)
         : primitive_base(id, inputs)
         , variable_info(variable_info)
         , concat_axis(concat_axis)
         , gather_axis(gather_axis)
-        , indirect(indirect) {}
+        , indirect(indirect)
+        , update_kv(update_kv) {}
 
     ov::op::util::VariableInfo variable_info;
     int64_t concat_axis = 0;
     int64_t gather_axis = 0;
     bool indirect = false;
-
+    bool update_kv = false;
     bool compressed = false;
     QuantizationAttributes quantization_attributes;
 
@@ -47,6 +49,7 @@ struct kv_cache : public primitive_base<kv_cache> {
         seed = hash_combine(seed, concat_axis);
         seed = hash_combine(seed, gather_axis);
         seed = hash_combine(seed, indirect);
+        seed = hash_combine(seed, update_kv);
         seed = hash_combine(seed, compressed);
         seed = hash_range(seed, quantization_attributes.scales_zp_output_order.begin(), quantization_attributes.scales_zp_output_order.end());
         seed = hash_range(seed, quantization_attributes.group_sizes.begin(), quantization_attributes.group_sizes.end());
@@ -69,6 +72,7 @@ struct kv_cache : public primitive_base<kv_cache> {
                concat_axis == rhs_casted.concat_axis &&
                gather_axis == rhs_casted.gather_axis &&
                indirect == rhs_casted.indirect &&
+               update_kv == rhs_casted.update_kv &&
                compressed == rhs_casted.compressed &&
                quantization_attributes.scales_zp_output_order == rhs_casted.quantization_attributes.scales_zp_output_order &&
                quantization_attributes.output_storage_type == rhs_casted.quantization_attributes.output_storage_type &&
@@ -88,6 +92,7 @@ struct kv_cache : public primitive_base<kv_cache> {
         ob << concat_axis;
         ob << gather_axis;
         ob << indirect;
+        ob << update_kv;
         ob << compressed;
         ob << make_data(&quantization_attributes.quantization_type, sizeof(quantization_attributes.quantization_type));
         ob << make_data(&quantization_attributes.quantization_dt, sizeof(quantization_attributes.quantization_dt));
@@ -110,6 +115,7 @@ struct kv_cache : public primitive_base<kv_cache> {
         ib >> concat_axis;
         ib >> gather_axis;
         ib >> indirect;
+        ib >> update_kv;
         ib >> compressed;
         ib >> make_data(&quantization_attributes.quantization_type, sizeof(quantization_attributes.quantization_type));
         ib >> make_data(&quantization_attributes.quantization_dt, sizeof(quantization_attributes.quantization_dt));