Skip to content

openvino: fix CPY writeback for recurrent state rollback #263

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wine99 merged 1 commit into from
Jul 28, 2026
Merged

openvino: fix CPY writeback for recurrent state rollback #263

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 42 additions & 15 deletions ggml/src/ggml-openvino/ggml-decoder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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 &&
Expand Down Expand Up @@ -569,21 +579,34 @@ std::pair<ModelParams, ComputeParams> 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];
}
}
}
Expand Down Expand Up @@ -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) {
Expand Down
12 changes: 11 additions & 1 deletion ggml/src/ggml-openvino/ggml-decoder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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<std::string, RsWriteback> 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 {
Expand Down
54 changes: 23 additions & 31 deletions ggml/src/ggml-openvino/openvino/op/cpy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int64_t>{1, 1, -1, output_shape[3].get_length()});

ov::Output<ov::Node> src;
ov::Output<ov::Node> begin;
ov::Output<ov::Node> 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::v1::Multiply>(
ov::op::v0::Constant::create(ov::element::i64, {1}, {ssm_state_size}), slot_len);
auto state_begin = std::make_shared<ov::op::v0::Negative>(state_rows);
auto state_part =
std::make_shared<ov::op::v8::Slice>(context.get_input(0), state_begin, int_max, one, axis);
auto feature = (int64_t) output_shape[3].get_length();
src = std::make_shared<ov::op::v1::Reshape>(
state_part,
ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{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<ov::op::v8::Slice>(context.get_input(0), src_begin, int_max, one, axis);
src = std::make_shared<ov::op::v1::Reshape>(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<ov::op::v8::Slice>(
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<ov::op::v1::Reshape>(
conv_state_last,
ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{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<ov::op::v1::Add>(
src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size}));
auto window = std::make_shared<ov::op::v8::Slice>(context.get_input(0), src_begin, src_end, one,
ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
src = std::make_shared<ov::op::v1::Reshape>(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<ov::op::v1::Add>(slot_idx, slot_len);
}

if (src.get_element_type() != context.get_output_type()) {
Expand Down
Loading