Enable -fa off

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

@@ -1,20 +1,15 @@
 #include "ggml-decoder.h"
 
-#include "ggml-backend-impl.h"
-#include "ggml-backend.h"
+#include "ggml-impl.h"
 #include "ggml-openvino-extra.h"
 #include "ggml-openvino.h"
 #include "ggml-quants.h"
 
-#include <ggml-impl.h>
-#include <ggml.h>
-
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <cstdlib>
-#include <execution>
 #include <fstream>
 #include <iomanip>
 #include <map>
@@ -30,12 +25,10 @@
 #include <openvino/op/convert.hpp>
 #include <openvino/op/parameter.hpp>
 #include <openvino/runtime/tensor.hpp>
-#include <optional>
 #include <ostream>
 #include <set>
 #include <stdexcept>
 #include <string>
-#include <unordered_map>
 #include <vector>
 
 GgmlOvDecoder::GgmlOvDecoder(ggml_cgraph * cgraph,
@@ -159,7 +152,7 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
             if (src->ne[2] * src->ne[3] == node->ne[1]) {
                 op_case = 5;
             }
-        } else if (src->ne[0] * src->ne[1] == node->ne[1]) {
+        } else if (src->ne[0] * src->ne[1] * src->ne[2] == node->ne[1]) {
             op_case = 3;
         } else if (src->ne[1] * src->ne[2] == node->ne[1]) {
             op_case = 6;
@@ -173,20 +166,40 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
             // kv cache tensor
             std::string src_name(node->view_src->name);
             int layer = extract_layer_from_name(src_name);
-            if (!is_swa_layer(layer)) {
-                op_case = 2;
+            if (ggml_is_contiguous(node->src[0])) {
+                // -  19: [    64,     8,   256,     1] VIEW            cache_k_l0 (view)             [ 2,   128,  1024, 1048576]
+                //         [   512,  1024,     1,     1]      0: NONE     cache_k_l0                    [ 2,  1024, 1048576, 1048576]
+                // -  20: [    64,   256,     8,     1] PERMUTE         cache_k_l0 (view) (permuted)  [ 2,  1024,   128, 1048576]
+                //         [    64,     8,   256,     1]      0: VIEW     cache_k_l0 (view)             [ 2,   128,  1024, 1048576]
+                if (!is_swa_layer(layer)) {
+                    op_case = 3;
+                } else {
+                    op_case = 4;
+                }
             } else {
-                op_case = 3;
+                // special case of cache v when `-fa off`
+                // -  17: [   256,     8,    64,     1] VIEW            cache_v_l0 (view)             [ 2, 131072,  2048, 1048576]
+                //         [   512,  1024,     1,     1]      0: NONE     cache_v_l0                   [ 2,  1024, 1048576, 1048576]
+                // -  18: [   256,    64,     8,     1] PERMUTE         cache_v_l0 (view) (permuted)  [ 2,  2048, 131072, 1048576]
+                //         [   256,     8,    64,     1]      0: VIEW     cache_v_l0 (view)            [ 2, 131072,  2048, 1048576]
+                if (!is_swa_layer(layer)) {
+                    op_case = 5;
+                } else {
+                    op_case = 6;
+                }
             }
         } else {
             // rope'ed query tensor
-            op_case = 4;
+            op_case = 2;
         }
         break;
     }
     case GGML_OP_MUL_MAT: {
         if (node->src[0]->op == GGML_OP_VIEW && node->src[1]->op == GGML_OP_VIEW) {
             op_case = 3;
+        } else if (node->src[1]->op == GGML_OP_SOFT_MAX) {
+            // In the case of `-fa off`, softmax is used, v_trans=true, the dynamic dim is ne[0] for cache_v
+            op_case = 2;
         }
         break;
     }
@@ -273,13 +286,20 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
     for (int i = 0; i < cgraph->n_nodes; i++) {
         auto * node = cgraph->nodes[i];
         std::string name = std::string(node->name);
-        if (node->op == GGML_OP_FLASH_ATTN_EXT) {
-            model_params.n_heads = node->src[0]->ne[2];
-            model_params.n_heads_kv = node->src[1]->ne[2];
-            model_params.head_size = node->src[0]->ne[0];
+        if (node->op == GGML_OP_FLASH_ATTN_EXT || node->op == GGML_OP_SOFT_MAX) {
             compute_params.input_len = node->src[0]->ne[1];
 
+            auto * q_perm = node->src[0];
             auto * cache_k_perm = node->src[1];
+            if (node->op == GGML_OP_SOFT_MAX) {
+                q_perm = node->src[0]->src[1];
+                cache_k_perm = node->src[0]->src[0];
+            }
+            model_params.head_size = cache_k_perm->ne[0];
+            model_params.n_heads_kv = cache_k_perm->ne[2];
+            model_params.n_heads = q_perm->ne[2];
+            compute_params.token_len_per_seq = q_perm->ne[1];
+
             if (cache_k_perm->op == GGML_OP_CPY) {
                 cache_k_perm = cache_k_perm->src[0];
             }
@@ -289,7 +309,11 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
 
             auto * cache_k = cache_k_view->src[0];
             int layer = extract_layer_from_name(cache_k->name);
+
             auto * mask = node->src[3];
+            if (node->op == GGML_OP_SOFT_MAX) {
+                mask = node->src[1];
+            }
             std::string mask_name(mask->name);
 
             model_params.kv_buffer_ctx_id = ggml_backend_openvino_buffer_get_ctx_id(cache_k->buffer);
@@ -306,7 +330,6 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
             size_t offset;
             memcpy(&offset, cache_k_view->op_params, sizeof(size_t));
             compute_params.seq_active_start = offset / seq_size;
-            compute_params.token_len_per_seq = node->ne[2];
 
             if (mask_name.find("swa") != std::string::npos) {
                 compute_params.attention_size_swa = mask->ne[0];
@@ -318,7 +341,6 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
                 compute_params.attention_size_swa = model_params.ctx_per_seq_swa;
                 compute_params.token_len_per_seq = 1;
             }
-            break;
         }
         // if the node op is TRANSPOSE and its input is PERMUTE and the source of the PERMUTE is VIEW, then get the attention size with the TRANSPOSE node ne[0] (in case no GGML_OP_FLASH_ATTN_EXT)
         if (node->op == GGML_OP_TRANSPOSE && node->src[0]->op == GGML_OP_PERMUTE &&
@@ -1286,4 +1308,4 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
             std::cout << std::endl;
         }
     }
-}
+}

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

@@ -1,6 +1,7 @@
 #pragma once
 
-#include "ggml-quants.h"
+#include "ggml-backend-impl.h"
+#include "ggml-backend.h"
 #include "ggml.h"
 #include "openvino/decoder.h"
 
@@ -9,7 +10,6 @@
 #include <map>
 #include <memory>
 #include <openvino/core/partial_shape.hpp>
-#include <optional>
 #include <vector>
 
 struct ModelParams {
@@ -239,15 +239,17 @@ class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
     }
 
     inline static bool is_inp_mask(const ggml_tensor * tensor, const ggml_tensor * op) {
-        return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]);
+        return op->op == GGML_OP_CPY || (op->op == GGML_OP_FLASH_ATTN_EXT && tensor == op->src[3]) ||
+               (op->op == GGML_OP_SOFT_MAX && tensor == op->src[1]);
     }
 
     inline static bool is_rope_freqs_weight(const ggml_tensor * tensor, const ggml_tensor * op) {
         return op->op == GGML_OP_ROPE && tensor == op->src[2];
     }
 
     inline static bool is_kvcache(const ggml_tensor * tensor, const ggml_tensor * op) {
-        return op->op == GGML_OP_SET_ROWS && op->src[2] == tensor;
+        return (op->op == GGML_OP_SET_ROWS && op->src[2] == tensor) ||
+               tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_ANY;
     }
 
     inline static bool is_kv_idx(const ggml_tensor * tensor, const ggml_tensor * op) {

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

@@ -34,23 +34,19 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
     auto q = std::make_shared<ov::op::v0::Convert>(q_f32, ov::element::f16);
     auto scale_node = std::make_shared<ov::op::v0::Constant>(ov::element::f16, ov::Shape{}, std::vector<float>{scale});
 
-    ov::Output<ov::Node> mask_sliced, res;
+    ov::Output<ov::Node> res;
+
+    // For stateful
     std::string mask_name = "KQ_mask_sliced";
     if (context.get_input_names()[3].find("swa") != std::string::npos) {
         mask_name = "KQ_mask_swa_sliced";
     }
     if (context.has_input(mask_name)) {
-        mask_sliced = context.get_input(mask_name);
-    } else {
-        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 two = ov::op::v0::Constant::create(ov::element::i64, {1}, {2});
-        auto token_len = get_dimensions(q, {2});
-        mask_sliced = std::make_shared<ov::op::v8::Slice>(mask, zero, token_len, one, two);
+        mask = context.get_input(mask_name);
     }
 
-    if (mask_sliced.get_element_type() != ov::element::f16) {
-        mask_sliced = std::make_shared<ov::op::v0::Convert>(mask_sliced, ov::element::f16);
+    if (mask.get_element_type() != ov::element::f16) {
+        mask = std::make_shared<ov::op::v0::Convert>(mask, ov::element::f16);
     }
 
     auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output<Node> kv) {
@@ -77,7 +73,7 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
     k = tile_kv(q_shape[1], k_shape[1], q_shape[3], k);
     v = tile_kv(q_shape[1], k_shape[1], q_shape[3], v);
 
-    auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, mask_sliced, scale_node, false);
+    auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, mask, scale_node, false);
     res = std::make_shared<ov::op::v1::Transpose>(sdpa,
                                                   ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
     res = std::make_shared<ov::op::v0::Convert>(res, ov::element::f32);

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

@@ -34,10 +34,7 @@ OutputVector translate_mulmat(const NodeContext & context) {
     ov::Output<ov::Node> A = context.get_input(1);
 
     bool transpose_b = true;
-    if (op_case == 2) {
-        B = B.get_node_shared_ptr()->input_value(0);
-        transpose_b = false;
-    } else if (op_case == 3) {
+    if (op_case == 3) {
         B = process_view_input(context, 0);
         A = process_view_input(context, 1);
     }
@@ -55,6 +52,7 @@ OutputVector translate_mulmat(const NodeContext & context) {
     auto batch_small = A_batch_larger ? B_batch : A_batch;
 
     Output<Node> Z = A_batch_larger ? B : A;
+    auto Z_shape = A_batch_larger ? B_shape : A_shape;
     int64_t factor = batch_large / batch_small;
     if (factor > 1 && batch_small > 1) {
         auto batch_large_node = ov::op::v0::Constant::create(ov::element::i64, {1}, std::vector<int64_t>{batch_large});
@@ -67,7 +65,11 @@ OutputVector translate_mulmat(const NodeContext & context) {
         auto broadcast_shape = ov::op::v0::Constant::create(
             ov::element::i64, {5}, {(int64_t) 1, (int64_t) 1, factor, (int64_t) 1, (int64_t) 1});
         auto new_Z_shape = ov::op::v0::Constant::create(ov::element::i64, {4},
-                                                        {(int64_t) 0, batch_large, (int64_t) -1, (int64_t) A_shape[3]});
+                                                        {(int64_t) 0, batch_large, (int64_t) -1, (int64_t) Z_shape[3]});
+        if (op_case == 2) {
+            new_Z_shape = ov::op::v0::Constant::create(ov::element::i64, {4},
+                                                       {(int64_t) 0, batch_large, (int64_t) Z_shape[2], (int64_t) -1});
+        }
 
         auto Z_broadcasted = std::make_shared<ov::op::v3::Broadcast>(Z_unsqueezed, broadcast_shape,
                                                                      ov::op::BroadcastType::BIDIRECTIONAL);

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

@@ -23,8 +23,11 @@ OutputVector translate_permute(const NodeContext & context) {
     num_inputs_check(context, 1, 1);
 
     int op_case = context.get_op_case();
-    FRONT_END_CHECK_IMPLEMENTED(op_case == 1 || op_case == 2 || op_case == 3 || op_case == 4,
-                                "Unsupported PERMUTE case");
+    FRONT_END_CHECK_IMPLEMENTED(op_case != 0, "Unsupported PERMUTE case");
+    // op_case 1 is trivial permute
+    // op_case 2 is to permute Q. It has a preceding VIEW that reshapes Q to restore the sequqence dimension
+    // op_case 3 4 it to permute KV cache in the default layout
+    // op_case 5 6 is to permute V cache when `-fa off`, where v_trans=true
 
     ov::Output<Node> res;
     auto src = context.get_input(0);
@@ -39,7 +42,7 @@ OutputVector translate_permute(const NodeContext & context) {
 
     if (op_case == 1 || context.is_stateful()) {
         res = std::make_shared<ov::op::v1::Transpose>(src, perm);
-    } else if (op_case == 4) {
+    } else if (op_case == 2) {
         auto output_shape = context.get_output_shape().to_shape();
         auto n_heads = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[1]});
         auto head_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[3]});
@@ -62,13 +65,17 @@ OutputVector translate_permute(const NodeContext & context) {
         auto output_shape = context.get_output_shape().to_shape();
         int64_t head_size = output_shape[3];
         int64_t n_heads = output_shape[1];
+        if (op_case == 5 || op_case == 6) {
+            head_size = output_shape[2];
+            n_heads = output_shape[1];
+        }
         int64_t ctx_per_seq = cache_shape[2].is_static() ? cache_shape[2].get_length() : -1;
         int64_t n_seq = cache_shape[1].get_length();
 
         Output<Node> attention_size;
         if (!context.has_input("attention_size")) {
             attention_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[2]});
-        } else if (op_case == 2) {
+        } else if (op_case == 3 || op_case == 5) {
             attention_size = context.get_input("attention_size");
         } else {
             attention_size = context.get_input("attention_size_swa");
@@ -88,18 +95,31 @@ OutputVector translate_permute(const NodeContext & context) {
             seq_active_end = ov::op::v0::Constant::create(ov::element::i64, {1}, {seq_active_end_val});
         }
 
-        // 1. reshape to [n_seq, ctx_per_seq, n_heads, head_size]
+        // 1. reshape to [n_seq, ctx_per_seq, n_heads, head_size] (for `-fa off` [n_seq, n_heads, head_size, ctx_per_seq])
         // 2. slice out the active sequences
         // 3. slice out the attention part in each sequence
-        // 4. permute
+        // 4. permute (skip for `-fa off`)
         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 src_reshaped = std::make_shared<ov::op::v1::Reshape>(
-            src, ov::op::v0::Constant::create(ov::element::i64, {4}, {n_seq, ctx_per_seq, n_heads, head_size}), false);
-        auto slice1 = std::make_shared<ov::op::v8::Slice>(src_reshaped, seq_active_start, seq_active_end, one, zero);
-        auto slice2 = std::make_shared<ov::op::v8::Slice>(slice1, zero, attention_size, one, one);
-        res = std::make_shared<ov::op::v1::Transpose>(slice2, perm);
+        if (op_case == 3 || op_case == 4) {
+            auto src_reshaped = std::make_shared<ov::op::v1::Reshape>(
+                src, ov::op::v0::Constant::create(ov::element::i64, {4}, {n_seq, ctx_per_seq, n_heads, head_size}),
+                false);
+            auto slice1 =
+                std::make_shared<ov::op::v8::Slice>(src_reshaped, seq_active_start, seq_active_end, one, zero);
+            auto slice2 = std::make_shared<ov::op::v8::Slice>(slice1, zero, attention_size, one, one);
+            res = std::make_shared<ov::op::v1::Transpose>(slice2, perm);
+        } else {
+            auto three = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
+            auto src_reshaped = std::make_shared<ov::op::v1::Reshape>(
+                src, ov::op::v0::Constant::create(ov::element::i64, {4}, {n_seq, n_heads, head_size, ctx_per_seq}),
+                false);
+            auto slice1 =
+                std::make_shared<ov::op::v8::Slice>(src_reshaped, seq_active_start, seq_active_end, one, zero);
+            auto slice2 = std::make_shared<ov::op::v8::Slice>(slice1, zero, attention_size, one, three);
+            res = slice2;
+        }
     }
     return rename_outputs_with_suffix({res}, context.get_name());
 }

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

@@ -10,7 +10,6 @@
 #include <openvino/op/concat.hpp>
 #include <openvino/op/constant.hpp>
 #include <openvino/op/reshape.hpp>
-#include <stdexcept>
 #include <vector>
 
 namespace ov {
@@ -47,7 +46,14 @@ OutputVector translate_reshape(const NodeContext & context) {
             std::vector<int64_t>{(int64_t) output_shape[0], (int64_t) output_shape[1], -1, (int64_t) output_shape[3]});
 
     } else if (op_case == 3) {
-        throw std::runtime_error("might be outdated RESHAPE case");
+        //  -  14: [     1,  1024,     1,     1] RESHAPE              Vcur-0 (reshaped) (reshaped)
+        //         [   512,     2,     1,     1]            0: RESHAPE     Vcur-0 (reshaped)
+        //  -  15: [     1, 524288,     1,     1] RESHAPE              cache_v_l0 (reshaped)
+        //         [   512,  1024,     1,     1]            0: NONE        cache_v_l0
+        //  -  16: [     1, 524288,     1,     1] SET_ROWS             cache_v_l0 (reshaped) (view)
+        //         [     1,  1024,     1,     1]            0: RESHAPE     Vcur-0 (reshaped) (reshaped)
+        //         [  1024,     1,     1,     1]            1: NONE        leaf_11
+        //         [     1, 524288,     1,     1]            2: RESHAPE     cache_v_l0 (reshaped)
         new_shape_node = ov::op::v0::Constant::create(
             ov::element::i64, {4}, std::vector<int64_t>{(int64_t) output_shape[0], (int64_t) output_shape[1], -1, 1});
 

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

@@ -34,14 +34,14 @@ OutputVector translate_set_rows(const NodeContext & context) {
 
     data = std::make_shared<ov::op::v0::Convert>(data, context.get_output_type());
 
-    auto dst_shape = context.get_output_shape().to_shape();
+    auto row_size = context.get_input_shape(2)[3].get_length();
 
     auto ind_squeezed =
         std::make_shared<ov::op::v0::Squeeze>(indices, ov::op::v0::Constant::create(ov::element::i64, {3}, {0, 1, 2}));
     auto data_reshaped = std::make_shared<ov::op::v1::Reshape>(
         data,
         ov::op::v0::Constant::create(ov::element::i64, {4},
-                                     {(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) dst_shape[3]}),
+                                     {(int64_t) 1, (int64_t) 1, (int64_t) -1, (int64_t) row_size}),
         false);
     auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{}, {2});