ggml-openvino: keep stateful execution with -fa off

Keep the V cache token-major in stateful mode instead of falling back to
stateless execution. SET_ROWS appends to the state, and PERMUTE reads it
back in the layout attention needs, so MUL_MAT does not transpose it.
Dropping a reshape in grouped attention lets OpenVINO fuse it into SDPA.

For padded V caches, slice after the transpose so the cache is read in
place, without a copy that is recompiled for every KV length.

This removes the stateless fallback for -fa off.

Author: osabnis

docs/backend/OPENVINO.md

@@ -801,7 +801,6 @@ build\ReleaseOV\bin\llama-simple.exe -m "C:\models\Llama-3.2-1B-Instruct-Q4_K_M.
 
 - `llama-bench`: requires `-fa 1` (flash-attention).
 - `llama-cli --context-shift`: stateless only (`GGML_OPENVINO_STATEFUL_EXECUTION=0`). In stateful mode the KV cache is owned by the OpenVINO model and cannot be shifted externally.
-- `-fa off`: stateless only. With `-fa off` the V cache is stored transposed, which the stateful KV cache does not support; with `GGML_OPENVINO_STATEFUL_EXECUTION=1` the context falls back to stateless execution and logs a warning.
 - `llama-server`: only one chat session/thread when `GGML_OPENVINO_STATEFUL_EXECUTION=1`.
 
 **GPU-specific**

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

@@ -452,6 +452,14 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
         }
         break;
     }
+    case GGML_OP_SET_ROWS: {
+        // `-fa off` writes the V cache transposed, through a [1, n] reshape of the cache: one element per row
+        if (node->src[2] != nullptr && node->src[2]->op == GGML_OP_RESHAPE && node->src[2]->ne[0] == 1 &&
+            is_kvcache(node->src[2]->src[0], nullptr)) {
+            op_case = 1;
+        }
+        break;
+    }
     case GGML_OP_MUL_MAT: {
         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

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

@@ -31,6 +31,8 @@ OutputVector translate_mulmat(const NodeContext & context) {
 
     ov::Output<ov::Node> B = process_view_input_new(context, 0);
     ov::Output<ov::Node> A = process_view_input_new(context, 1);
+    // the stateful V cache with `-fa off` is not transposed (see permute)
+    const bool transpose_b = B.get_node()->get_rt_info().count(GGML_OV_V_TOKEN_MAJOR) == 0;
 
     if (A.get_element_type() != B.get_element_type()) {
         B = std::make_shared<ov::op::v0::Convert>(B, context.get_input_type(1));
@@ -55,9 +57,19 @@ OutputVector translate_mulmat(const NodeContext & context) {
         //   A [n, batch_large, rows, k] -> [n, batch_small, factor * rows, k] (a pure reshape: the group is adjacent)
         //   A x B^T [n, batch_small, factor * rows, cols] -> [n, batch_large, rows, cols]
         // The zeros copy the input dimensions (special_zero), which may be dynamic.
-        auto A_grouped = std::make_shared<ov::op::v1::Reshape>(
+        ov::Output<Node> A_grouped = std::make_shared<ov::op::v1::Reshape>(
             A, ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, batch_small, -1, 0}), true);
-        ov::Output<Node> res = std::make_shared<ov::op::v0::MatMul>(A_grouped, B, false, true);
+        // A is often a Reshape of the grouped scores (see softmax): take them directly. With one Reshape less, the
+        // OpenVINO SDPAFusion can match the attention
+        const auto * A_reshape = ov::as_type<ov::op::v1::Reshape>(A.get_node());
+        const auto * A_target =
+            A_reshape != nullptr ? ov::as_type<ov::op::v0::Constant>(A_reshape->get_input_node_ptr(1)) : nullptr;
+        if (A_target != nullptr && A_reshape->get_special_zero() &&
+            A_target->cast_vector<int64_t>() == std::vector<int64_t>{0, batch_large, -1, 0} &&
+            A_reshape->get_input_partial_shape(0)[1] == batch_small) {
+            A_grouped = A_reshape->input_value(0);
+        }
+        ov::Output<Node> res = std::make_shared<ov::op::v0::MatMul>(A_grouped, B, false, transpose_b);
         res = std::make_shared<ov::op::v1::Reshape>(
             res, ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, batch_large, -1, 0}),
             true);
@@ -96,7 +108,6 @@ OutputVector translate_mulmat(const NodeContext & context) {
         A = Z;
     }
 
-    bool transpose_b = true;
     ov::Output<Node> res = std::make_shared<ov::op::v0::MatMul>(A, B, false, transpose_b);
 
     const auto output_type = context.get_output_type();

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

@@ -76,8 +76,14 @@ OutputVector translate_pad(const NodeContext & context) {
         return rename_outputs_with_suffix({res}, context.get_name());
     }
 
-    const std::vector<int64_t> pads_begin = {pads[6], pads[4], pads[2], pads[0]};
-    const std::vector<int64_t> pads_end = {pads[7], pads[5], pads[3], pads[1]};
+    std::vector<int64_t> pads_begin = {pads[6], pads[4], pads[2], pads[0]};
+    std::vector<int64_t> pads_end = {pads[7], pads[5], pads[3], pads[1]};
+    // stateful execution drops the leading size-1 batch dim
+    const auto input_rank = input.get_partial_shape().rank();
+    if (input_rank.is_static() && input_rank.get_length() == 3 && pads_begin[0] == 0 && pads_end[0] == 0) {
+        pads_begin.erase(pads_begin.begin());
+        pads_end.erase(pads_end.begin());
+    }
 
     auto pads_begin_node = ov::op::v0::Constant::create(ov::element::i64, {pads_begin.size()}, pads_begin);
     auto pads_end_node = ov::op::v0::Constant::create(ov::element::i64, {pads_end.size()}, pads_end);

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

@@ -46,7 +46,40 @@ OutputVector translate_permute(const NodeContext & context) {
                 static_cast<int64_t>(perm_values.size() - 1 - input_axis);
         }
     }
-    if (op_case == 1 || context.is_stateful()) {
+    if (context.is_stateful() && (op_case == 5 || op_case == 6)) {
+        // `-fa off` reads V transposed. The stateful V cache is token-major, [1, seq, n_heads_kv, head_size]
+        const auto output_shape = context.get_output_shape().to_shape();
+        const int64_t n_heads = output_shape[1];
+        const int64_t head_size = output_shape[2];
+        const auto & state_shape = src.get_partial_shape();
+        const bool single_head = n_heads == 1 && state_shape[2].get_length() == 1;
+        if (!single_head && state_shape[2].get_length() * state_shape[3].get_length() != n_heads * head_size) {
+            // llama.cpp pads the V cache rows to the widest layer: keep this layer's rows
+            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 rows = std::make_shared<ov::op::v1::Reshape>(
+                src, ov::op::v0::Constant::create(ov::element::i64, {3}, std::vector<int64_t>{0, 0, -1}), true);
+            auto sliced = std::make_shared<ov::op::v8::Slice>(
+                rows, zero, ov::op::v0::Constant::create(ov::element::i64, {1}, {n_heads * head_size}), one,
+                ov::op::v0::Constant::create(ov::element::i64, {1}, {2}));
+            src = std::make_shared<ov::op::v1::Reshape>(
+                sliced,
+                ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, 0, n_heads, head_size}),
+                true);
+        }
+        // keep V as [1, n_heads_kv, kv, head_size]: MUL_MAT then multiplies without transposing it
+        res = std::make_shared<ov::op::v1::Transpose>(
+            src, ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{0, 2, 1, 3}));
+        if (single_head && state_shape[3].get_length() != head_size) {
+            // the padded rows, as above. After the Transpose, which then cancels the state relayout, so the Slice reads
+            // the KV cache in place: between the two they become copies of the cache, recompiled for every length
+            res = std::make_shared<ov::op::v8::Slice>(res, ov::op::v0::Constant::create(ov::element::i64, {1}, {0}),
+                                                      ov::op::v0::Constant::create(ov::element::i64, {1}, {head_size}),
+                                                      ov::op::v0::Constant::create(ov::element::i64, {1}, {1}),
+                                                      ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
+        }
+        res.get_node()->get_rt_info()[GGML_OV_V_TOKEN_MAJOR] = true;
+    } else if (op_case == 1 || context.is_stateful()) {
         // The stateful path carries hidden-state tensors in a rank-3 layout (the
         // leading batch dim is dropped, e.g. Gemma4's per-layer-embedding path). The
         // perm above is rank-4; when the actual input is rank-3, drop the batch axis

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

@@ -55,11 +55,13 @@ OutputVector translate_set_rows(const NodeContext & context) {
     Output<Node> res;
     if (context.is_stateful()) {
         int concat_axis = 1;
-        // the stateful cache appends whole tokens; with `-fa off` llama.cpp writes the V cache transposed,
-        // one element per row
+        if (context.get_op_case() == 1) {
+            // `-fa off` writes V transposed through a [1, n] reshape of the cache. The state keeps V token-major
+            // like K, and the new values are already in token order, so append them to the state itself
+            dst = dst.get_node()->input_value(0);
+        }
         FRONT_END_OP_CONVERSION_CHECK(dst.get_partial_shape()[2].is_static() && dst.get_partial_shape()[3].is_static(),
-                                      "Stateful execution needs flash attention (-fa on): with -fa off the V cache "
-                                      "is written transposed");
+                                      "Unsupported KV cache layout for stateful execution");
         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>(

ggml/src/ggml-openvino/openvino/utils.h

@@ -86,6 +86,9 @@ inline bool is_kvcache_4d_shape(const ov::PartialShape & shape) {
 // stacked per B matrix
 constexpr const char * GGML_OV_GROUPED_ROWS = "ggml_ov_grouped_rows";
 
+// rt_info key on the stateful V cache read with `-fa off` (permute): V is [.., kv, head_size], not transposed
+constexpr const char * GGML_OV_V_TOKEN_MAJOR = "ggml_ov_v_token_major";
+
 namespace op {
 template <typename T> OutputVector translate_1to1_match_2_inputs(const NodeContext & context) {
     num_inputs_check(context, 2, 2);

ggml/src/ggml-openvino/utils.cpp

@@ -1066,12 +1066,6 @@ enum ggml_status ov_graph_compute_dynamic(ggml_cgraph * cgraph, const std::share
     const int64_t graph_param_compute_time = ggml_time_us() - graph_param_compute_start_time;
 
     const bool cache_enabled = !model_is_splitted && !cache_disabled;
-    if (r_ctx->stateful && m_params.v_trans) {
-        GGML_LOG_WARN(
-            "ggml-openvino: disabling stateful execution for this context: the graph writes a transposed V cache "
-            "(flash attention is off), which the stateful KV cache does not support\n");
-        r_ctx->stateful = false;
-    }
     const bool stateful_recurrent = r_ctx->stateful && m_params.state_size >= 0;
     const bool stateful_kv_only = r_ctx->stateful && !stateful_recurrent;
     if (r_ctx->stateful &&