diff --git a/ggml/src/ggml-openvino/ggml-decoder.cpp b/ggml/src/ggml-openvino/ggml-decoder.cpp index 16f059380759..130668958db0 100644 --- a/ggml/src/ggml-openvino/ggml-decoder.cpp +++ b/ggml/src/ggml-openvino/ggml-decoder.cpp @@ -115,6 +115,16 @@ bool is_same_shape(const ggml_tensor * a, const ggml_tensor * b) { bool is_conv_states_all_tensor(const ggml_tensor * tensor) { return tensor != nullptr && strncmp(tensor->name, "conv_states_all", strlen("conv_states_all")) == 0; } + +// CPY writing the tail of conv_input (the concat of the previous conv state and the new tokens) +// back into a slot block of the recurrent state cache. Detected structurally because the rollback +// variant (cparams.n_rs_seq > 0) emits one such CPY per snapshot slot without naming them. +bool is_conv_state_writeback(const ggml_tensor * node) { + return node->op == GGML_OP_CPY && node->view_src != nullptr && GgmlOvDecoder::is_kvcache(node->view_src, nullptr) && + node->src[0] != nullptr && node->src[0]->op == GGML_OP_VIEW && node->src[0]->src[0] != nullptr && + node->src[0]->src[0]->op == GGML_OP_CONCAT && node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && + node->src[1]->view_src == node->view_src; +} } // namespace static std::string get_tensor_ov_name(const ggml_cgraph * cgraph, const ggml_tensor * tensor) { @@ -353,7 +363,7 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { if (node->src[0]->op == GGML_OP_VIEW) { if (node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET) { op_case = 1; - } else if (std::string(node->src[0]->name).find("conv_state_last") == 0) { + } else if (is_conv_state_writeback(node)) { op_case = 2; break; } else if (is_conv_states_all_tensor(node->view_src) && node->src[1] != nullptr && @@ -569,21 +579,34 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr compute_params.cache_rs_reset_len = ggml_nelements(node) / node->view_src->ne[0]; compute_params.cache_rs_reset_idx = node->src[0]->view_offs / node->view_src->ne[0]; } - // Capture the active-slot block of the recurrent state reorder (inp->s_copy). The active - // sequences occupy a contiguous slot block [idx, idx+len) of the state cache; read both from - // the active conv/gdn state writeback destination view (idx = head, len = n_seqs). + // Capture the destination slot block of every recurrent state cache writeback, plus the + // conv_input window the conv state writeback copies. The active sequences occupy a + // contiguous slot block [begin, begin + n_seqs) of the cache; the block and the window move + // with the batch, so they are fed to the cached model as runtime inputs. if (node->op == GGML_OP_CPY && node->view_src != nullptr && is_kvcache(node->view_src, nullptr) && - node->src[0]->op == GGML_OP_VIEW && node->src[1] != nullptr) { - const bool is_conv = std::string(node->src[0]->name).find("conv_state_last") == 0; - const bool is_gdn = node->src[0]->src[0] != nullptr && node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET; - if (is_conv || is_gdn) { - const ggml_tensor * dest_view = node->src[1]; - const ggml_tensor * cache = node->view_src; - const size_t row_bytes = cache->ne[0] * ggml_type_size(cache->type); - if (row_bytes > 0) { - compute_params.s_copy_active_slot_idx = (int) (dest_view->view_offs / row_bytes); - compute_params.s_copy_active_slot_len = (int) dest_view->ne[1]; + node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src == node->view_src) { + const bool is_conv = is_conv_state_writeback(node); + const bool is_gdn = node->src[0]->op == GGML_OP_VIEW && node->src[0]->src[0] != nullptr && + node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET; + const bool is_extra = node->src[0]->op == GGML_OP_GET_ROWS; + + const ggml_tensor * dest_view = node->src[1]; + const ggml_tensor * cache = node->view_src; + const size_t row_bytes = cache->ne[0] * ggml_type_size(cache->type); + if (row_bytes > 0 && (is_conv || is_gdn || is_extra)) { + ComputeParams::RsWriteback writeback; + writeback.slot_begin = (int) (dest_view->view_offs / row_bytes); + if (is_conv) { + // conv_input column the copied window starts at + writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[0]); + } else if (is_gdn) { + // first row of the state part of the gated-delta-net output + writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[1]); } + compute_params.rs_writebacks[get_tensor_ov_name(cgraph, node)] = writeback; + } + if (is_conv || is_gdn) { + compute_params.s_copy_active_slot_len = (int) dest_view->ne[1]; } } } @@ -730,9 +753,13 @@ void GgmlOvDecoder::add_extra_inputs() { } if (m_compute_params.s_copy_active_slot_len != -1) { - create_1d_input("s_copy_active_slot_idx", m_compute_params.s_copy_active_slot_idx); create_1d_input("s_copy_active_slot_len", m_compute_params.s_copy_active_slot_len); } + + for (const auto & [node_name, writeback] : m_compute_params.rs_writebacks) { + create_1d_input("rs_slot_begin_" + node_name, writeback.slot_begin); + create_1d_input("rs_src_begin_" + node_name, writeback.src_begin); + } } bool GgmlOvDecoder::node_is_used_as_src(const int node_idx) { diff --git a/ggml/src/ggml-openvino/ggml-decoder.h b/ggml/src/ggml-openvino/ggml-decoder.h index 1372d3ebdfb0..b4979fcae29f 100644 --- a/ggml/src/ggml-openvino/ggml-decoder.h +++ b/ggml/src/ggml-openvino/ggml-decoder.h @@ -60,7 +60,6 @@ struct ComputeParams { // 5: [ 18432, 1, 1, 1] SCALE cache_r_l0 (reshaped) (view) (view) // [ 18432, 1, 1, 1] 0: VIEW cache_r_l0 (reshaped) (view) - int s_copy_active_slot_idx = -1; int s_copy_active_slot_len = -1; // SSM/DeltaNet models otionally reorder slots of state cache, to make the active slots contiguous // leaf_5 is the inp->s_copy in llama-graph.cpp, eg if there are 8 slots in total and slot 3 and 7 @@ -80,6 +79,17 @@ struct ComputeParams { // 11: [ 18432, 0, 1, 1] CPY cache_r_l0 (view) (copy of ) // [ 18432, 0, 1, 1] 0: GET_ROWS node_9 // [ 18432, 0, 1, 1] 1: VIEW cache_r_l0 (view) + + struct RsWriteback { + int slot_begin = 0; // first cache slot written by the CPY + int src_begin = 0; // where the copied data starts in the source tensor (in rows of it) + }; + + std::map rs_writebacks; + // Offsets of the state cache writeback CPY nodes, keyed by node name. They change with the + // batch (kv head, active sequence count, token count) and, with rollback enabled + // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot, each snapshot + // taking a different conv_input window. Passed to the cached model as runtime inputs. }; class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder { diff --git a/ggml/src/ggml-openvino/openvino/op/cpy.cpp b/ggml/src/ggml-openvino/openvino/op/cpy.cpp index be9772efdd0c..8d3234bba6a1 100644 --- a/ggml/src/ggml-openvino/openvino/op/cpy.cpp +++ b/ggml/src/ggml-openvino/openvino/op/cpy.cpp @@ -60,50 +60,42 @@ OutputVector translate_cpy(const NodeContext & context) { return rename_outputs_with_suffix({res}, context.get_name()); } - // Recurrent state cache writeback with a dynamic active-slot block (inp->s_copy reorder). - // The active sequences occupy a contiguous slot block [idx, idx+len) of the state cache; write - // the new rows into that block while preserving the rest, so the result is the full updated - // cache. op_case 1: gated-delta-net state, op_case 2: conv state, op_case 3: defrag remainder. - const bool slice_assign = context.has_input("s_copy_active_slot_len") && !context.is_stateful() && - (op_case == 1 || op_case == 2 || op_case == 3); + // Recurrent state cache writeback into a slot block of the cache. Where the block starts and + // where the copied data starts in the source are runtime inputs, so the cached model works for + // any kv head, active sequence count and token count. The result is the full updated cache. + // op_case 1: gated-delta-net state, op_case 2: conv state, op_case 3: defrag remainder. + const std::string slot_begin_name = "rs_slot_begin_" + context.get_name(); + const bool slice_assign = + context.has_input(slot_begin_name) && !context.is_stateful() && (op_case >= 1 && op_case <= 3); if (slice_assign) { const int64_t slot_axis = 2; - auto slot_idx = context.get_input("s_copy_active_slot_idx"); - auto slot_len = context.get_input("s_copy_active_slot_len"); auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); auto int_max = ov::op::v0::Constant::create(ov::element::i64, {1}, {INT_MAX}); auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {slot_axis}); + auto feature = ov::op::v0::Constant::create(ov::element::i64, {4}, + std::vector{1, 1, -1, output_shape[3].get_length()}); ov::Output src; - ov::Output begin; + ov::Output begin = context.get_input(slot_begin_name); if (op_case == 1) { - // GDN packs [attn | new_state]; the state is the last ssm_state_size * n_seqs rows. - int ssm_state_size = context.get_ssm_state_size(); - auto state_rows = std::make_shared( - ov::op::v0::Constant::create(ov::element::i64, {1}, {ssm_state_size}), slot_len); - auto state_begin = std::make_shared(state_rows); - auto state_part = - std::make_shared(context.get_input(0), state_begin, int_max, one, axis); - auto feature = (int64_t) output_shape[3].get_length(); - src = std::make_shared( - state_part, - ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector{1, 1, -1, feature}), false); - begin = slot_idx; + // GDN packs [attn | state snapshots]; the state part runs from src_begin to the end. + auto src_begin = context.get_input("rs_src_begin_" + context.get_name()); + auto state_part = std::make_shared(context.get_input(0), src_begin, int_max, one, axis); + src = std::make_shared(state_part, feature, false); } else if (op_case == 2) { - auto cache_r_size = (int64_t) input_shape[3].get_length(); - auto conv_state_last = std::make_shared( - context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {1}, {-cache_r_size}), int_max, - one, ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); - auto feature = (int64_t) output_shape[3].get_length(); - src = std::make_shared( - conv_state_last, - ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector{1, 1, -1, feature}), false); - begin = slot_idx; + // conv_input is [previous conv state | new tokens]; copy the conv_kernel_size - 1 wide + // window starting at src_begin, which is the snapshot this writeback corresponds to. + auto window_size = (int64_t) input_shape[3].get_length(); + auto src_begin = context.get_input("rs_src_begin_" + context.get_name()); + auto src_end = std::make_shared( + src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size})); + auto window = std::make_shared(context.get_input(0), src_begin, src_end, one, + ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); + src = std::make_shared(window, feature, false); } else { // op_case 3: gathered remainder rows already have the cache slot layout [1, 1, extra, feature] src = context.get_input(0); - begin = std::make_shared(slot_idx, slot_len); } if (src.get_element_type() != context.get_output_type()) {