ggml-openvino: support gemma-2 soft-capped attention

Gemma-2 caps its attention scores as softcap * tanh(scores / softcap)
before the softmax. FLASH_ATTN_EXT now applies the cap on the manual
attention path, as SDPA has no soft-capping, and the decoder looks
through the SCALE, TANH, SCALE chain of -fa off when it matches
attention.

One helper finds the K cache of an attention for both the
sliding-window pre-pass and the main pass. The pre-pass copy did not
see through the chain, so the sliding-window layers used the
full-attention sizes once the context exceeded n_swa + n_ubatch.

Soft-capped attention keeps the flat KV cache layout: with the 4-D
layout the GPU plugin fails in its oneDNN path.

Author: osabnis

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

@@ -762,6 +762,16 @@ static int get_swa_window_from_mask(const ggml_tensor * mask) {
     return window > 0 ? (int) window : -1;
 }
 
+// The KQ product that the SOFT_MAX of an attention reads: directly, or through the logit soft-capping of gemma-2
+// (SCALE, TANH, SCALE)
+static const ggml_tensor * softmax_scores_source(const ggml_tensor * scores) {
+    while (scores != nullptr && (scores->op == GGML_OP_SCALE ||
+                                 (scores->op == GGML_OP_UNARY && ggml_get_unary_op(scores) == GGML_UNARY_OP_TANH))) {
+        scores = scores->src[0];
+    }
+    return scores;
+}
+
 std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgraph * cgraph, bool is_static) {
     ModelParams model_params;
     ComputeParams compute_params;
@@ -804,12 +814,12 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
                 break;
             }
             break;
-        case GGML_OP_SOFT_MAX:
-            // case 2: node op is SOFT_MAX, src 0 not null & op is MUL_MAT & the src 0 of MUL_MAT is PERMUTE & the permuted tensor src is the view of cache k
-            if (node->src[0] != nullptr && node->src[1] != nullptr && node->src[0]->op == GGML_OP_MUL_MAT &&
-                node->src[0]->src[0] != nullptr && node->src[0]->src[1] != nullptr &&
-                node->src[0]->src[0]->op == GGML_OP_PERMUTE && node->src[0]->src[0]->src[0] != nullptr &&
-                node->src[0]->src[0]->src[0]->op == GGML_OP_VIEW) {
+        case GGML_OP_SOFT_MAX: {
+            // case 2: node op is SOFT_MAX, its scores come from a MUL_MAT (see softmax_scores_source) & the src 0 of MUL_MAT is PERMUTE & the permuted tensor src is the view of cache k
+            const ggml_tensor * kq = softmax_scores_source(node->src[0]);
+            if (kq != nullptr && node->src[1] != nullptr && kq->op == GGML_OP_MUL_MAT && kq->src[0] != nullptr &&
+                kq->src[1] != nullptr && kq->src[0]->op == GGML_OP_PERMUTE && kq->src[0]->src[0] != nullptr &&
+                kq->src[0]->src[0]->op == GGML_OP_VIEW) {
                 return 2;
             }
             // case 3: node op is SOFT_MAX, src 0 not null & op is ADD & the src 0 of ADD is MUL_MAT & the src 0 of MUL_MAT is PERMUTE
@@ -819,6 +829,7 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
                 return 3;
             }
             break;
+        }
         default:
             break;
         }
@@ -841,6 +852,22 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
         }
     };
 
+    // The PERMUTE of the K cache view an attention node reads, for the same src layouts
+    auto get_attention_cache_k_permute = [&get_attention_pattern_case](const ggml_tensor * node) -> ggml_tensor * {
+        switch (get_attention_pattern_case(node)) {
+        case 0:
+            return node->src[1];
+        case 1:
+            return node->src[1]->src[0];
+        case 2:
+            return softmax_scores_source(node->src[0])->src[0];
+        case 3:
+            return node->src[0]->src[0]->src[0];
+        default:
+            return nullptr;
+        }
+    };
+
     // Pre-pass: classify sliding-window vs full-attention layers.
     //
     // An interleaved-SWA model keeps two KV caches and two attention masks, and hands each layer
@@ -874,14 +901,7 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
             if (mask == nullptr) {
                 continue;
             }
-            const ggml_tensor * cache_k_permute = nullptr;
-            switch (get_attention_pattern_case(cgraph->nodes[i])) {
-            case 0:  cache_k_permute = cgraph->nodes[i]->src[1];                     break;
-            case 1:  cache_k_permute = cgraph->nodes[i]->src[1]->src[0];             break;
-            case 2:  cache_k_permute = cgraph->nodes[i]->src[0]->src[0];             break;
-            default: cache_k_permute = cgraph->nodes[i]->src[0]->src[0]->src[0];     break;
-            }
-            const ggml_tensor * cache_k_view = cache_k_permute->src[0];
+            const ggml_tensor * cache_k_view = get_attention_cache_k_permute(cgraph->nodes[i])->src[0];
             if (cache_k_view->op != GGML_OP_VIEW) {
                 continue;
             }
@@ -950,26 +970,23 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
             node->src[2]->ne[0] == 1 && is_kvcache(node->src[2]->src[0], nullptr)) {
             model_params.v_trans = true;
         }
+        if (node->op == GGML_OP_FLASH_ATTN_EXT) {
+            float logit_softcap = 0.0f;
+            memcpy(&logit_softcap, (const float *) node->op_params + 2, sizeof(float));
+            model_params.softcap_attn = model_params.softcap_attn || logit_softcap != 0.0f;
+        }
         const int attention_pattern_case = get_attention_pattern_case(node);
         if (attention_pattern_case != -1) {
-            ggml_tensor * cache_k_permute = nullptr;
+            ggml_tensor * cache_k_permute = get_attention_cache_k_permute(node);
             ggml_tensor * mask = nullptr;
 
             switch (attention_pattern_case) {
             case 0:
-                cache_k_permute = node->src[1];
-                mask = node->src[3];
-                break;
             case 1:
-                cache_k_permute = node->src[1]->src[0];
                 mask = node->src[3];
                 break;
             case 2:
-                cache_k_permute = node->src[0]->src[0];
-                mask = node->src[1];
-                break;
             case 3:
-                cache_k_permute = node->src[0]->src[0]->src[0];
                 mask = node->src[1];
                 break;
             case 4:
@@ -1593,8 +1610,9 @@ bool GgmlOvDecoder::is_kvcache_4d(const ggml_tensor * cache) const {
         return false;
     }
     // `-fa off` (transposed V cache): keep the [1, 1, ctx, n_heads_kv * head_size] layout. The 4-D K cache breaks the
-    // MUL_MAT attention of the Phi-3 models (Phi-3.5 mini, Phi-4 mini).
-    if (m_model_params.v_trans) {
+    // MUL_MAT attention of the Phi-3 models (Phi-3.5 mini, Phi-4 mini), and the soft-capped manual attention of
+    // gemma-2 with `-fa on` ("could not construct a memory descriptor using strides" in the GPU plugin's oneDNN path).
+    if (m_model_params.v_trans || m_model_params.softcap_attn) {
         return false;
     }
     if (cache->op != GGML_OP_NONE || cache->view_src != nullptr || !is_kvcache(cache, nullptr)) {

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

@@ -37,6 +37,8 @@ struct ModelParams {
     bool is_cacheless_attn = false;
     // the graph writes a transposed V cache (`-fa off`: one value element per row, tokens along the row)
     bool v_trans = false;
+    // the graph has FLASH_ATTN_EXT with logit soft-capping (gemma-2)
+    bool softcap_attn = false;
     std::vector<int> swa_layers;
     // The sliding-window mask tensor, identified in compute_llm_params() by grouping attention
     // layers on the mask they consume. Only used to tell the two masks apart when naming OV

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

@@ -1646,8 +1646,10 @@ static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
         if (max_bias > 0) {
             return {false, "FLASH_ATTN_EXT with max_bias > 0 (max_bias=" + std::to_string(max_bias) + ") is not supported"};
         }
-        if (logit_softcap != 0) {
-            return {false, "FLASH_ATTN_EXT with logit_softcap != 0 (logit_softcap=" + std::to_string(logit_softcap) + ") is not supported"};
+        // soft-capping (gemma-2) is applied on the manual attention path of translate_flash_attn_ext, which needs one
+        // head size for Q, K and V
+        if (logit_softcap != 0 && (op->src[1]->ne[0] != op->src[0]->ne[0] || op->src[2]->ne[0] != op->src[0]->ne[0])) {
+            return {false, "FLASH_ATTN_EXT with logit_softcap != 0 needs one head size for Q, K and V"};
         }
         break;
     }

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

@@ -19,6 +19,7 @@
 #include <openvino/op/scatter_update.hpp>
 #include <openvino/op/slice.hpp>
 #include <openvino/op/softmax.hpp>
+#include <openvino/op/tanh.hpp>
 #include <openvino/op/tile.hpp>
 #include <openvino/op/transpose.hpp>
 #include <openvino/op/unsqueeze.hpp>
@@ -136,7 +137,7 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
     float * params = reinterpret_cast<float *>(context.get_output_op_params());
     float scale = params[0];
     // float max_bias      = params[1];
-    // float logit_softcap = params[2];
+    const float logit_softcap = params[2];
 
     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});
@@ -183,8 +184,14 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
     // The manual path reshapes Q, K, V and the output with one head size; models whose K/V head sizes differ
     // from Q's (e.g. MLA) take the SDPA path.
     const bool same_head_size = k_shape[3] == q_shape[3] && v_shape[3] == q_shape[3];
+    // SDPA has no logit soft-capping (gemma-2), so soft-capped attention always takes the manual path, also in
+    // stateful execution
     const bool use_manual_gqa_attention =
-        manual_gqa_enabled && factor > 1 && num_heads_kv > 1 && same_head_size && !context.is_stateful();
+        (logit_softcap != 0.0f || (manual_gqa_enabled && factor > 1 && num_heads_kv > 1 && !context.is_stateful())) &&
+        same_head_size;
+    FRONT_END_OP_CONVERSION_CHECK(logit_softcap == 0.0f || use_manual_gqa_attention,
+                                  "FLASH_ATTN_EXT with logit soft-capping is not supported with different K/V head "
+                                  "sizes");
 
     if (use_manual_gqa_attention) {
         // Q, K, V arrive as [B, n_heads(_kv), S, head_size], where B is the active
@@ -205,7 +212,18 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
 
         // QK^T -> [B, num_heads_kv, factor * S_q, S_k]
         auto qk = std::make_shared<ov::op::v0::MatMul>(q_r, k, /*tA=*/false, /*tB=*/true);
-        ov::Output<ov::Node> probs = std::make_shared<ov::op::v1::Multiply>(qk, scale_node);
+        ov::Output<ov::Node> probs;
+        if (logit_softcap != 0.0f) {
+            // as ggml: softcap * tanh(qk * scale / softcap)
+            probs = std::make_shared<ov::op::v1::Multiply>(
+                qk,
+                ov::op::v0::Constant::create(ov::element::f16, ov::Shape{}, std::vector<float>{scale / logit_softcap}));
+            probs = std::make_shared<ov::op::v1::Multiply>(
+                std::make_shared<ov::op::v0::Tanh>(probs),
+                ov::op::v0::Constant::create(ov::element::f16, ov::Shape{}, std::vector<float>{logit_softcap}));
+        } else {
+            probs = std::make_shared<ov::op::v1::Multiply>(qk, scale_node);
+        }
 
         if (has_mask) {
             // Mask arrives as [B, 1, S_q, S_k]. Row j * S_q + s of the scores belongs to query position s, so

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

@@ -121,7 +121,9 @@ OutputVector translate_permute(const NodeContext & context) {
 
         Output<Node> attention_size;
         if (!context.has_input("attention_size")) {
-            attention_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[2]});
+            // the attended tokens are the last axis of the transposed V cache (`-fa off`)
+            const size_t tokens_axis = (op_case == 5 || op_case == 6) ? 3 : 2;
+            attention_size = ov::op::v0::Constant::create(ov::element::i64, {1}, {output_shape[tokens_axis]});
         } else if (op_case == 3 || op_case == 5) {
             attention_size = context.get_input("attention_size");
         } else {