Function planner_set_compact_kv_row

src/models/ace_step/planner.cpp:410–440

    return {output, k, v};
}

core::TensorValue planner_set_compact_kv_row(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & cache,
    const core::TensorValue & row,
    const core::TensorValue & cache_slot) {
    core::validate_rank_between(cache, 4, 4, "cache");
    core::validate_shape(
        row,
        core::TensorShape::from_dims({cache.shape.dims[0], 1, cache.shape.dims[2], cache.shape.dims[3]}),
        "row");
    const int64_t batch = cache.shape.dims[0];
    core::validate_shape(cache_slot, core::TensorShape::from_dims({batch}), "cache_slot");
    if (batch == 1) {
        const modules::FastKVSetRowsModule set_rows;
        return set_rows.build(ctx, cache, row, cache_slot);
    }

    const int64_t steps = cache.shape.dims[1];
    const int64_t row_elems = cache.shape.dims[2] * cache.shape.dims[3];
    auto flat_cache = core::reshape_tensor(
        ctx,
        cache,
        core::TensorShape::from_dims({batch * steps, row_elems}));
    auto flat_row = core::reshape_tensor(
        ctx,
        ensure_planner_contiguous(ctx, row),
        core::TensorShape::from_dims({batch, row_elems}));
    ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, cache_slot.tensor);
    auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type);
    return core::reshape_tensor(ctx, flat_updated, cache.shape);
}

modules::QwenDecoderLayerOutputs planner_decoder_layer_with_static_cache_tail_batched(