NPU unify PD. Unify dynamic and static dims

Author: wine99

ggml/src/ggml-openvino/ggml-decoder.cpp

@@ -27,7 +27,6 @@
 #include <openvino/op/constant.hpp>
 #include <openvino/op/convert.hpp>
 #include <openvino/op/parameter.hpp>
-#include <openvino/op/unsqueeze.hpp>
 #include <openvino/runtime/tensor.hpp>
 #include <optional>
 #include <ostream>
@@ -39,7 +38,6 @@
 GgmlOvDecoder::GgmlOvDecoder(ggml_tensor * node,
                              ggml_cgraph * cgraph,
                              bool is_static,
-                             bool is_first_token,
                              int context_size,
                              int context_size_swa,
                              int num_heads,
@@ -55,25 +53,24 @@ GgmlOvDecoder::GgmlOvDecoder(ggml_tensor * node,
     m_num_heads(num_heads),
     m_num_heads_kv(num_heads_kv),
     m_head_size(head_size),
-    m_is_static(is_static),
-    m_is_first_token(is_first_token) {
+    m_is_static(is_static) {
     set_input_output(node);
 }
 
 GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph,
                              std::map<std::string, std::shared_ptr<ov::Node>> & model_weights,
-                             bool is_static,
-                             bool is_first_token) :
+                             bool is_static) :
     m_cgraph(cgraph),
     m_op_name(m_node ? std::string(m_node->name) : ""),
     m_model_weights(model_weights),
-    m_is_static(is_static),
-    m_is_first_token(is_first_token) {
-    if (is_first_token && getenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS")) {
+    m_is_static(is_static) {
+    if (auto * env = getenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS"); env && std::string(env) != "0") {
+        unsetenv("GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS");
         print_tensor_address_map(cgraph);
     }
 
     set_llm_params();
+    validate_cgraph();
 
     for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) {
         auto * cur_node = cgraph->nodes[node_n];
@@ -300,41 +297,39 @@ void GgmlOvDecoder::set_llm_params() {
     }
 }
 
+void GgmlOvDecoder::validate_cgraph() const {
+    if (m_is_static && m_input_len != 1) {
+        throw std::runtime_error("Static graph (NPU) must have input_len == 1, but got " + std::to_string(m_input_len) +
+                                 ", try set -ub 1");
+    }
+}
+
 ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * src) const {
     auto name = std::string(src->name);
     ov::PartialShape input_shape;
-    if (name == "inp_tokens" || name == "inp_pos") {
-        if (m_is_static) {
-            if (m_is_first_token) {
-                input_shape = ov::PartialShape{1, 1, m_context_size};
-            } else {
-                input_shape = ov::PartialShape{1, 1, 1};
-            }
-        } else {
-            input_shape = ov::PartialShape{1, 1, -1};
-        }
-    } else if (name == "inp_out_ids" && !m_is_static) {
-        input_shape = ov::PartialShape{1, 1, -1};
+
+    if (name == "inp_tokens" || name == "inp_pos" || name == "inp_out_ids") {
+        input_shape = ov::PartialShape{1, 1, m_is_static ? 1 : -1};
+
     } else if (name.find("KQ_mask") == 0) {
         if (m_is_static) {
-            if (m_is_first_token) {
-                input_shape = ov::PartialShape{1, m_context_size, m_context_size};
-            } else {
-                input_shape = ov::PartialShape{1, 1, m_context_size};
-            }
+            input_shape = ov::PartialShape{1, 1, m_context_size};
         } else {
             input_shape = ov::PartialShape{1, -1, -1};
         }
+
     } else if (name.find("cache_") == 0) {
+        auto past_token_len = -1;
         if (m_is_static) {
             int layer = extract_layer_from_name(name);
             bool is_swa = is_swa_layer(layer);
-            input_shape = ov::PartialShape{is_swa ? m_context_size_swa : m_context_size, m_num_heads_kv, m_head_size};
-        } else {
-            input_shape = ov::PartialShape{1, -1, m_num_heads_kv, m_head_size};
+            past_token_len = is_swa ? m_context_size_swa : m_context_size;
         }
+        input_shape = ov::PartialShape{past_token_len, m_num_heads_kv, m_head_size};
+
     } else if (const auto * op = get_tensor_used_op(src); op && op->op == GGML_OP_SET_ROWS) {
         input_shape = ov::PartialShape{1, 1, m_is_static ? 1 : -1};
+
     } else if (src->op == GGML_OP_VIEW) {
         // This case is added to make test-backend-ops work
         input_shape = ov::PartialShape{get_shape(src->view_src)};
@@ -748,9 +743,8 @@ int32_t * GgmlOvDecoder::get_output_op_params(const std::string & name) const {
 
 void GgmlOvDecoder::visit_subgraph(std::function<void(std::shared_ptr<GgmlDecoder>)> node_visitor) const {
     for (const auto & node : m_nodes) {
-        auto decoder =
-            std::make_shared<GgmlOvDecoder>(node, m_cgraph, m_is_static, m_is_first_token, m_context_size,
-                                            m_context_size_swa, m_num_heads, m_num_heads_kv, m_head_size, m_swa_layers);
+        auto decoder = std::make_shared<GgmlOvDecoder>(node, m_cgraph, m_is_static, m_context_size, m_context_size_swa,
+                                                       m_num_heads, m_num_heads_kv, m_head_size, m_swa_layers);
         node_visitor(decoder);
     }
 }

ggml/src/ggml-openvino/ggml-decoder.h

@@ -16,14 +16,12 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
     // Graph decoder
     GgmlOvDecoder(ggml_cgraph * cgraph,
                   std::map<std::string, std::shared_ptr<ov::Node>> & model_weights,
-                  bool is_static,
-                  bool is_first_token);
+                  bool is_static);
 
     // Node decoder, called in GgmlOvDecoder::visit_subgraph
     GgmlOvDecoder(ggml_tensor * node,
                   ggml_cgraph * cgraph,
                   bool is_static,
-                  bool is_first_token,
                   int context_size,
                   int context_size_swa,
                   int num_heads,
@@ -129,8 +127,6 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
 
     virtual bool is_static() const override { return m_is_static; }
 
-    virtual bool is_first_token() const override { return m_is_first_token; }
-
     ov::PartialShape get_graph_input_shape(const ggml_tensor * src) const;
 
     static void dump_cgraph(const ggml_cgraph * cgraph, std::string & filename);
@@ -157,6 +153,7 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
 
     // set context_size, num_heads, etc
     void set_llm_params();
+    void validate_cgraph() const;
 
     ggml_cgraph * m_cgraph = nullptr;
     ggml_tensor * m_node = nullptr;
@@ -185,7 +182,6 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
     int32_t * m_rope_params;
     std::vector<std::string> m_kv_names;
     bool m_is_static = false;
-    bool m_is_first_token;
 };
 
 void print_tensor_address_map(const ggml_cgraph * cgraph);

ggml/src/ggml-openvino/openvino/decoder.hpp

@@ -65,7 +65,6 @@ class GgmlDecoder : public DecoderBase {
     virtual std::map<std::string, std::string> get_kv_param_res_names() const = 0;
 
     virtual bool is_static() const = 0;
-    virtual bool is_first_token() const = 0;
     virtual int get_context_size() const = 0;
     virtual int get_context_size_swa() const = 0;
     virtual int is_swa_layer(int layer) const = 0;

ggml/src/ggml-openvino/openvino/node_context.hpp

@@ -97,12 +97,7 @@ class NodeContext : public frontend::NodeContext {
     int get_op_case() const {
         return m_decoder->get_op_case();
     }
-    bool is_static() const {
-        return m_decoder->is_static();
-    }
-    bool is_first_token() const {
-        return m_decoder->is_first_token();
-    }
+    bool is_static() const { return m_decoder->is_static(); }
 
 private:
     std::shared_ptr<GgmlDecoder> m_decoder;

ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp

@@ -53,35 +53,22 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
 
         auto stop = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{token_len, kv_len}, 0);
         mask_sliced = std::make_shared<ov::op::v8::Slice>(mask, zero_2d, stop, one_2d, axes);
-        mask_sliced = std::make_shared<ov::op::v0::Unsqueeze>(mask_sliced, zero_1d);
     }
 
     if (mask_sliced.get_element_type() != ov::element::f16) {
         mask_sliced = std::make_shared<ov::op::v0::Convert>(mask_sliced, ov::element::f16);
     }
 
-    auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output<Node> kv,
-                       bool is_static) {
+    auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output<Node> kv) {
         int64_t factor = num_heads / num_heads_kv;
         if (factor > 1) {
             ov::Output<ov::Node> kv_broadcast_shape, kv_unsqueezed, new_kv_shape;
-            if (is_static) {
-                auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {1});
-                kv_unsqueezed = std::make_shared<ov::op::v0::Unsqueeze>(kv, unsqueeze_axes);
+            auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {1});
+            kv_unsqueezed = std::make_shared<ov::op::v0::Unsqueeze>(kv, unsqueeze_axes);
 
-                kv_broadcast_shape =
-                    ov::op::v0::Constant::create(ov::element::i64, {4}, {num_heads_kv, factor, (int64_t) 1, head_size});
-                new_kv_shape =
-                    ov::op::v0::Constant::create(ov::element::i64, {3}, {num_heads, (int64_t) -1, head_size});
-            } else {
-                auto unsqueeze_axes = ov::op::v0::Constant::create(ov::element::i64, Shape{}, {2});
-                kv_unsqueezed = std::make_shared<ov::op::v0::Unsqueeze>(kv, unsqueeze_axes);
-
-                kv_broadcast_shape = ov::op::v0::Constant::create(
-                    ov::element::i64, {5}, {(int64_t) 1, num_heads_kv, factor, (int64_t) 1, head_size});
-                new_kv_shape = ov::op::v0::Constant::create(ov::element::i64, {4},
-                                                            {(int64_t) 1, num_heads, (int64_t) -1, head_size});
-            }
+            kv_broadcast_shape =
+                ov::op::v0::Constant::create(ov::element::i64, {4}, {num_heads_kv, factor, (int64_t) 1, head_size});
+            new_kv_shape = ov::op::v0::Constant::create(ov::element::i64, {3}, {num_heads, (int64_t) -1, head_size});
 
             kv = std::make_shared<ov::op::v3::Broadcast>(kv_unsqueezed, kv_broadcast_shape,
                                                          ov::op::BroadcastType::BIDIRECTIONAL);
@@ -92,18 +79,12 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
 
     auto q_shape = context.get_input_shape(0).to_shape();
     auto k_shape = context.get_input_shape(1).to_shape();
-    k = tile_kv(q_shape[0], k_shape[0], q_shape[2], k, context.is_static());
-    v = tile_kv(q_shape[0], k_shape[0], q_shape[2], v, context.is_static());
+    k = tile_kv(q_shape[0], k_shape[0], q_shape[2], k);
+    v = tile_kv(q_shape[0], k_shape[0], q_shape[2], v);
 
     auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, mask_sliced, scale_node, false);
-    auto sdpa_f32 = std::make_shared<ov::op::v0::Convert>(sdpa, ov::element::f32);
-    if (context.is_static()) {
-        res = std::make_shared<ov::op::v1::Transpose>(sdpa_f32,
-                                                      ov::op::v0::Constant::create(ov::element::i64, {3}, {1, 0, 2}));
-    } else {
-        res = std::make_shared<ov::op::v1::Transpose>(
-            sdpa_f32, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
-    }
+    res = std::make_shared<ov::op::v1::Transpose>(sdpa, ov::op::v0::Constant::create(ov::element::i64, {3}, {1, 0, 2}));
+    res = std::make_shared<ov::op::v0::Convert>(res, ov::element::f32);
     return rename_outputs_with_suffix({res}, context.get_name());
 }
 

ggml/src/ggml-openvino/openvino/op/permute.cpp

@@ -26,40 +26,8 @@ OutputVector translate_permute(const NodeContext & context) {
     ov::Output<Node> res;
     auto zero = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
 
-    if (op_case == 1) {
-        if (context.is_static()) {
-            res = std::make_shared<ov::op::v1::Transpose>(
-                context.get_input(0), ov::op::v0::Constant::create(ov::element::i64, {3}, {1, 0, 2}));
-        } else {
-            auto src = context.get_input(0);
-            if (src.get_partial_shape().rank() == 3) {
-                src = std::make_shared<ov::op::v0::Unsqueeze>(src, zero);
-            }
-            res = std::make_shared<ov::op::v1::Transpose>(
-                src, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
-        }
-    } else {
-        auto src = context.get_input(0);
-        auto one = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
-
-        if (context.is_static()) {
-            auto src_shape_ = context.get_input_shape(0).to_shape();
-            std::vector<int64_t> src_shape(src_shape_.begin(), src_shape_.end());
-            auto src_reshaped = std::make_shared<ov::op::v1::Reshape>(
-                src,
-                ov::op::v0::Constant::create(ov::element::i64, {3},
-                                             std::vector<int64_t>{-1, src_shape[1], src_shape[2]}),
-                false);
-            res = std::make_shared<ov::op::v1::Transpose>(
-                src_reshaped, ov::op::v0::Constant::create(ov::element::i64, {3}, {1, 0, 2}));
-        } else {
-            if (src.get_partial_shape().rank() == 3) {
-                src = std::make_shared<ov::op::v0::Unsqueeze>(src, zero);
-            }
-            res = std::make_shared<ov::op::v1::Transpose>(
-                src, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
-        }
-    }
+    auto src = context.get_input(0);
+    res = std::make_shared<ov::op::v1::Transpose>(src, ov::op::v0::Constant::create(ov::element::i64, {3}, {1, 0, 2}));
     return rename_outputs_with_suffix({res}, context.get_name());
 }
 

ggml/src/ggml-openvino/openvino/op/rope.cpp

@@ -84,10 +84,6 @@ OutputVector translate_rope(const NodeContext & context) {
                                                               ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
         auto stack = std::make_shared<ov::op::v0::Concat>(OutputVector{first_half, second_half}, 3);
         res = std::make_shared<ov::op::v1::Reshape>(stack, std::make_shared<ov::op::v0::ShapeOf>(data_node), false);
-        if (!(context.is_static())) {
-            res =
-                std::make_shared<ov::op::v0::Unsqueeze>(res, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}));
-        }
     } else if (mode == ROPE_TYPE_NEOX) {
         auto data_split = std::make_shared<ov::op::v1::Split>(
             data_node, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {2}), 2);

ggml/src/ggml-openvino/openvino/op/set_rows.cpp

@@ -33,10 +33,6 @@ OutputVector translate_set_rows(const NodeContext & context) {
     auto dst_shape = context.get_output_shape(0).to_shape();
     FRONT_END_OP_CONVERSION_CHECK(dst_shape[0] == 1, "Unsupported shape in SET_ROWS");
 
-    if (context.is_static() && context.is_first_token()) {
-        return rename_outputs_with_suffix({data}, context.get_name());
-    }
-
     auto indices = context.get_input(1);
     auto dst = context.get_input(context.get_output_name());
 
@@ -54,13 +50,11 @@ OutputVector translate_set_rows(const NodeContext & context) {
         auto updated = std::make_shared<ov::op::v3::ScatterUpdate>(dst_reshaped, indices_reshaped, data_reshaped, zero);
         res = std::make_shared<ov::op::v1::Reshape>(updated, std::make_shared<ov::op::v0::ShapeOf>(dst), false);
     } else {
-        assert(dst.get_partial_shape().rank() == 4 && dst.get_partial_shape()[2].is_static() &&
-               dst.get_partial_shape()[3].is_static());
+        int64_t dim1 = dst.get_partial_shape()[1].get_length();
         int64_t dim2 = dst.get_partial_shape()[2].get_length();
-        int64_t dim3 = dst.get_partial_shape()[3].get_length();
         data = std::make_shared<ov::op::v1::Reshape>(
-            data, ov::op::v0::Constant::create(ov::element::i64, {4}, {(int64_t) 1, (int64_t) -1, dim2, dim3}), false);
-        res = std::make_shared<ov::op::v0::Concat>(OutputVector{dst, data}, 1);
+            data, ov::op::v0::Constant::create(ov::element::i64, {3}, {(int64_t) -1, dim1, dim2}), false);
+        res = std::make_shared<ov::op::v0::Concat>(OutputVector{dst, data}, 0);
     }
     return rename_outputs_with_suffix({res}, context.get_name());
 }

ggml/src/ggml-openvino/openvino/translate_session.cpp

@@ -27,7 +27,6 @@
 #include <openvino/op/squeeze.hpp>
 #include <openvino/op/strided_slice.hpp>
 #include <openvino/op/transpose.hpp>
-#include <openvino/op/unsqueeze.hpp>
 #include <openvino/pass/constant_folding.hpp>
 #include <openvino/pass/make_stateful.hpp>
 
@@ -112,7 +111,6 @@ void add_sliced_mask(TensorMap & tensor_map, GgmlDecoder & ggml_model_decoder) {
                 auto stop = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{token_len, kv_len}, 0);
 
                 mask_sliced = std::make_shared<ov::op::v8::Slice>(mask, zero_2d, stop, one_2d, axes);
-                mask_sliced = std::make_shared<ov::op::v0::Unsqueeze>(mask_sliced, zero_1d);
                 mask_sliced = std::make_shared<ov::op::v0::Convert>(mask_sliced, ov::element::f16);
                 mask_sliced->set_friendly_name(sliced_name);
             }