char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Zijun YuOpenVINO: Qwen3.5, memory optimization, and test-recurrent-state-rollback (#26952)aee56b3ab

char/aliceai-t5-moe
4.5 KiB95 linesraw
1#include "../node_context.h"
2#include "../op_table.h"
3#include "../utils.h"
4
5#include <cassert>
6#include <cstdint>
7#include <memory>
8#include <openvino/core/node.hpp>
9#include <openvino/core/node_output.hpp>
10#include <openvino/frontend/exception.hpp>
11#include <openvino/op/broadcast.hpp>
12#include <openvino/op/concat.hpp>
13#include <openvino/op/constant.hpp>
14#include <openvino/op/convert.hpp>
15#include <openvino/op/gather.hpp>
16#include <openvino/op/reshape.hpp>
17#include <openvino/op/scatter_elements_update.hpp>
18#include <openvino/op/scatter_update.hpp>
19#include <openvino/op/shape_of.hpp>
20#include <openvino/op/slice.hpp>
21#include <openvino/op/squeeze.hpp>
22#include <openvino/op/transpose.hpp>
23#include <vector>
24
25namespace ov {
26namespace frontend {
27namespace ggml {
28namespace op {
29
30OutputVector translate_set_rows(const NodeContext & context) {
31    num_inputs_check(context, 3, 3);
32
33    auto data = process_view_input_new(context, 0);
34    auto indices = process_view_input_new(context, 1);
35    auto dst = process_view_input_new(context, 2);
36
37    data = std::make_shared<ov::op::v0::Convert>(data, context.get_output_type());
38
39    const auto indices_shape = context.get_input_shape(1);
40    const bool multidim_indices = indices_shape.rank().is_static() &&
41                                  indices_shape.rank().get_length() == 4 &&
42                                  ((indices_shape[1].is_static() && indices_shape[1].get_length() > 1) ||
43                                   (indices_shape[2].is_static() && indices_shape[2].get_length() > 1));
44
45    auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {2});
46
47    Output<Node> res;
48    if (context.is_stateful()) {
49        int concat_axis = 1;
50        int64_t dim2 = dst.get_partial_shape()[2].get_length();
51        int64_t dim3 = dst.get_partial_shape()[3].get_length();
52        data = std::make_shared<ov::op::v1::Reshape>(
53            data, ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 1, (int64_t) -1, dim2, dim3}), false);
54        res = std::make_shared<ov::op::v0::Concat>(OutputVector{dst, data}, concat_axis);
55    } else if (multidim_indices) {
56        auto updates_shape = std::make_shared<ov::op::v3::ShapeOf>(data, ov::element::i64);
57
58        auto indices_rank3 = std::make_shared<ov::op::v0::Squeeze>(
59            indices, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}));
60        auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
61        auto indices_rank4_shape = std::make_shared<ov::op::v0::Concat>(OutputVector{get_dimensions(updates_shape, {0, 1, 2}), one}, 0);
62        auto indices_rank4 = std::make_shared<ov::op::v1::Reshape>(indices_rank3, indices_rank4_shape, false);
63        auto broadcasted_indices = std::make_shared<ov::op::v3::Broadcast>(indices_rank4, updates_shape);
64
65        res = std::make_shared<ov::op::v3::ScatterElementsUpdate>(dst, broadcasted_indices, data, axes);
66    } else {
67        auto row_size = context.get_input_shape(2)[3].get_length();
68        auto ind_squeezed = std::make_shared<ov::op::v0::Squeeze>(
69            indices, ov::op::v0::Constant::create(ov::element::i64, {3}, {0, 1, 2}));
70        auto data_reshaped = std::make_shared<ov::op::v1::Reshape>(
71            data,
72            ov::op::v0::Constant::create(ov::element::i64, {4},
73                                         {(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) row_size}),
74            false);
75        res = std::make_shared<ov::op::v3::ScatterUpdate>(dst, ind_squeezed, data_reshaped, axes);
76    }
77
78    auto dst_reshape = std::dynamic_pointer_cast<ov::op::v1::Reshape>(dst.get_node_shared_ptr());
79    if (!multidim_indices && dst_reshape) {
80        // Fix the case of multiple sequences, reshape back to original shape [1, n_seq, ctx_per_seq, emb]
81        // ctx_per_seq is not fixed due to llama-bench compatibility
82        auto dst_shape_partial = dst_reshape->get_input_partial_shape(0);
83        std::vector<int64_t> dst_shape = {dst_shape_partial[0].get_length(), dst_shape_partial[1].get_length(),
84                                          dst_shape_partial[2].is_static() ? dst_shape_partial[2].get_length() : -1,
85                                          dst_shape_partial[3].get_length()};
86        res = std::make_shared<ov::op::v1::Reshape>(res, ov::op::v0::Constant::create(ov::element::i64, {4}, dst_shape),
87                                                    false);
88    }
89    return rename_outputs_with_suffix({res}, context.get_name());
90}
91
92}  // namespace op
93}  // namespace ggml
94}  // namespace frontend
95}  // namespace ov