ggml-zdnn: introduce flash attention op (#30512)

* ggml-zdnn: add scratch buffer pool

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* ggml-zdnn: implement fattn, ditch scratch slots in favour of scratch buffer pool

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: add note about fattn

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

---------

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

docs/backend/zDNN.md

@@ -2,6 +2,7 @@
 
 > [!WARNING]
 > **Note:** zDNN is **not** the same as ZenDNN.
+>
 > - **zDNN** (this page): IBM's Deep Neural Network acceleration library for IBM Z & LinuxONE Mainframes
 > - **ZenDNN**: AMD's deep learning library for AMD EPYC CPUs ([see ZenDNN documentation](ZenDNN.md))
 
@@ -46,6 +47,14 @@ The IBM zDNN backend has the following CMake options that control the behaviour
 | --------------------- | ------------- | -------------------------------------------------------------------------------------- |
 | `GGML_ZDNN_HIGH_PREC` | `0`           | Set to `1` to requantize Q8_0 weights in smaller groups, for lower loss at lower speed |
 
+## Performance Tunables
+
+| Flag Option             | Default Value | Description                                            |
+| ----------------------- | ------------- | ------------------------------------------------------ |
+| `--flash-attn` or `-fa` | `auto`        | Set to `on` to run attention on the NNPA (Recommended) |
+
+Note: Set `-fa on` explicitly. Flash Attention fuses each layer's attention (the KQ matrix multiplication, softmax and KQV matrix multiplication) into one `FLASH_ATTN_EXT` operation, which zDNN runs on the NNPA instead of the CPU. The speedup grows with the context length. With `auto`, llama.cpp currently turns Flash Attention off, because the operation runs on a different device than the layer. zDNN runs it when the KV cache is F16 (the default), the model uses no attention sinks, ALiBi or logit softcapping, and there is one KV cache stream (one sequence, or `--kv-unified`). Otherwise it falls back to the CPU.
+
 ## 1. Install zDNN Library
 
 Note: A minimum of libzdnn 1.1.2 is required for the IBM zDNN backend to link correctly.

ggml/src/ggml-zdnn/common.hpp

@@ -6,6 +6,8 @@
 
 #include "zdnn.h"
 
+#include "scratch.hpp"
+
 #include <algorithm>
 #include <cstdlib>
 #include <future>
@@ -14,9 +16,8 @@
 #include <utility>
 #include <vector>
 
-#define GGML_ZDNN_NAME          "zDNN"
-#define GGML_ZDNN_VERSION       ZDNN_VERNUM
-#define GGML_ZDNN_SCRATCH_SLOTS 3  // int8 inputs, DFloat16 partial results, partial results for grouped matmul
+#define GGML_ZDNN_NAME    "zDNN"
+#define GGML_ZDNN_VERSION ZDNN_VERNUM
 
 #define ZDNN_CHECK(stmt)                \
     do {                                \
@@ -68,44 +69,100 @@ struct ggml_backend_zdnn_context {
     // zero bias for stacked quantized matmuls, keyed by (n_stacks, ne0)
     std::map<std::pair<int64_t, int64_t>, std::unique_ptr<ggml_backend_zdnn_buffer>> zero_bias;
 
-    std::vector<int8_t> scratch_q;
-    std::vector<float>  scratch_f32;
-
-    void * scratch_ztensor[GGML_ZDNN_SCRATCH_SLOTS]      = {};  // 4K aligned
-    size_t scratch_ztensor_size[GGML_ZDNN_SCRATCH_SLOTS] = {};
+    ggml_zdnn_pool pool;  // scratch memory of the ops
 };
 
-template <typename T>
-static inline T * ggml_zdnn_resize_scratch_max(std::vector<T> & scratch, int64_t n) {
-    if ((int64_t)scratch.size() < n) {
-        scratch.resize(n);
-    }
-
-    return scratch.data();
-}
-
 //
 // Scratch buffers
 //
 
-static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context * ctx,
-                                                  zdnn_tensor_desc          * pre_tfm_desc,
-                                                  zdnn_tensor_desc          * tfm_desc,
-                                                  zdnn_ztensor              * ztensor,
-                                                  int32_t                     slot) {
+// the ztensor uses the memory of scratch, so scratch must outlive it
+static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context     * ctx,
+                                                  zdnn_tensor_desc              * pre_tfm_desc,
+                                                  zdnn_tensor_desc              * tfm_desc,
+                                                  zdnn_ztensor                  * ztensor,
+                                                  ggml_zdnn_pool_alloc<uint8_t> & scratch) {
 
     const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
 
-    if (ctx->scratch_ztensor_size[slot] < size) {
-        free(ctx->scratch_ztensor[slot]);
-        ctx->scratch_ztensor[slot]      = std::aligned_alloc(4096, GGML_PAD(size, 4096));
-        ctx->scratch_ztensor_size[slot] = GGML_PAD(size, 4096);
-        GGML_ASSERT(ctx->scratch_ztensor[slot] != nullptr);
+    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
+    ztensor->buffer      = scratch.alloc(ctx->pool, size);
+    ztensor->buffer_size = size;
+}
+
+// same as ggml_zdnn_init_scratch_ztensor, for a DLFLOAT16 ztensor
+static inline void ggml_zdnn_init_scratch_ztensor_dlf16(ggml_backend_zdnn_context     * ctx,
+                                                        zdnn_tensor_desc              * pre_tfm_desc,
+                                                        zdnn_tensor_desc              * tfm_desc,
+                                                        zdnn_ztensor                  * ztensor,
+                                                        ggml_zdnn_pool_alloc<uint8_t> & scratch) {
+
+    const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
+
+    zdnn_init_ztensor(pre_tfm_desc, tfm_desc, ztensor);
+    ztensor->buffer      = scratch.alloc(ctx->pool, size);
+    ztensor->buffer_size = size;
+}
+
+//
+// Stacked ztensors
+//
+
+// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
+// the view descs must outlive the view
+static inline zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
+                                                zdnn_tensor_desc * tfm_desc,
+                                          const zdnn_ztensor     & src,
+                                                int64_t            n_stacks,
+                                                int64_t            s0,
+                                                int64_t            n) {
+
+    GGML_ASSERT(src.buffer_size % n_stacks == 0);
+    const uint64_t stack_bytes = src.buffer_size / n_stacks;
+
+    *pre_tfm_desc = *src.pre_transformed_desc;
+    *tfm_desc     = *src.transformed_desc;
+
+    switch (pre_tfm_desc->layout) {
+        case ZDNN_3DS:
+            {
+                pre_tfm_desc->dim3 = n;
+            } break;
+        case ZDNN_2DS:
+            {
+                pre_tfm_desc->dim2 = n;
+            } break;
+        default:
+            GGML_ABORT("%s: unsupported stacked layout", __func__);
     }
 
-    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
-    ztensor->buffer      = ctx->scratch_ztensor[slot];
-    ztensor->buffer_size = size;
+    tfm_desc->dim4 = n;
+    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
+
+    zdnn_ztensor view = src;
+    view.pre_transformed_desc = pre_tfm_desc;
+    view.transformed_desc     = tfm_desc;
+    view.buffer               = (char *) src.buffer + s0*stack_bytes;
+    view.buffer_size          = n * stack_bytes;
+    return view;
+}
+
+// zero bias of n_stacks stacks, created once per shape as it never changes
+static inline const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
+    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
+
+    if (!bias) {
+        bias = std::make_unique<ggml_backend_zdnn_buffer>();
+
+        std::vector<float> zeros(n_stacks * ne0, 0.0f);
+
+        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
+        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
+        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
+    }
+
+    return bias->ztensor;
 }
 
 //

ggml/src/ggml-zdnn/fattn.cpp

@@ -0,0 +1,255 @@
+#include "ggml.h"
+#include "fattn.hpp"
+
+#include <algorithm>
+#include <cstring>
+
+// DLFLOAT16 has no infinity, so masked scores use the lowest FP16 value, exp of which is still 0
+static inline float ggml_zdnn_fattn_mask_value(ggml_fp16_t m) {
+    return std::max(GGML_FP16_TO_FP32(m), -65504.0f);
+}
+
+bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op) {
+    const ggml_tensor * q    = op->src[0];
+    const ggml_tensor * k    = op->src[1];
+    const ggml_tensor * v    = op->src[2];
+    const ggml_tensor * mask = op->src[3];
+
+    float max_bias      = 0.0f;
+    float logit_softcap = 0.0f;
+    memcpy(&max_bias,      (const float *) op->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float));
+
+    // sinks, alibi, softcap and sparse k/v sets are not implemented
+    if (op->src[4] != nullptr || max_bias != 0.0f || logit_softcap != 0.0f || ggml_get_op_params_i32(op, 4) != 0) {
+        return false;
+    }
+
+    if (op->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) {
+        return false;
+    }
+
+    if (mask != nullptr && (mask->type != GGML_TYPE_F16 || mask->ne[2] != 1 || mask->ne[3] != 1)) {
+        return false;
+    }
+
+    if (q->ne[3] != 1 || k->ne[3] != 1 || v->ne[3] != 1 || k->ne[2] != v->ne[2] || q->ne[2] % k->ne[2] != 0) {
+        return false;
+    }
+
+    // rows are gathered with their strides, but each row has to be contiguous
+    if (q->nb[0] != sizeof(float) || k->nb[0] != sizeof(ggml_fp16_t) || v->nb[0] != sizeof(ggml_fp16_t)) {
+        return false;
+    }
+
+    const int64_t max_dim = ctx_dev->max_size;
+    const int64_t group   = q->ne[2] / k->ne[2];
+
+    return q->ne[0] <= max_dim && v->ne[0] <= max_dim && k->ne[1] <= max_dim &&
+           q->ne[2] <= max_dim && group * q->ne[1] <= max_dim;
+}
+
+void ggml_zdnn_flash_attn_ext(ggml_backend_zdnn_context * ctx, ggml_tensor * dst) {
+    const ggml_tensor * q    = dst->src[0];
+    const ggml_tensor * k    = dst->src[1];
+    const ggml_tensor * v    = dst->src[2];
+    const ggml_tensor * mask = dst->src[3];
+
+    const int64_t d_k       = q->ne[0];
+    const int64_t n_q       = q->ne[1];
+    const int64_t n_head    = q->ne[2];
+    const int64_t n_kv      = k->ne[1];
+    const int64_t n_head_kv = k->ne[2];
+    const int64_t d_v       = v->ne[0];
+    const int64_t group     = n_head / n_head_kv;
+
+    float scale = 1.0f;
+    memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
+
+    const int32_t n_threads = n_q*n_kv >= 65536 ? ctx->n_threads : 1;
+
+    // one query runs all heads in one call with the mask row as the K*Q bias, more queries run per head and add the mask
+    const bool single = n_q == 1;
+
+    const int64_t n_mask = mask ? (single ? n_head_kv*n_kv : n_q*n_kv) : 0;
+
+    ggml_zdnn_pool_alloc<float>       q_host_alloc(ctx->pool, n_head*n_q*d_k);
+    ggml_zdnn_pool_alloc<float>       o_host_alloc(ctx->pool, n_head*n_q*d_v);
+    ggml_zdnn_pool_alloc<float>       mask_host_alloc(ctx->pool, n_mask);
+    ggml_zdnn_pool_alloc<ggml_fp16_t> k_host_alloc(ctx->pool, n_head_kv*n_kv*d_k);
+    ggml_zdnn_pool_alloc<ggml_fp16_t> v_host_alloc(ctx->pool, n_head_kv*n_kv*d_v);
+
+    float       * q_host    = q_host_alloc.get();
+    float       * o_host    = o_host_alloc.get();
+    float       * mask_host = mask_host_alloc.get();
+    ggml_fp16_t * k_host    = k_host_alloc.get();
+    ggml_fp16_t * v_host    = v_host_alloc.get();
+
+    // q is [n_head][n_q][d_k] with the scale folded in, k and v are [n_head_kv][n_kv][d]
+    ggml_zdnn_parallel_for(n_threads, n_head, [&](int64_t h_start, int64_t h_end) {
+        for (int64_t h = h_start; h < h_end; h++) {
+            for (int64_t iq = 0; iq < n_q; iq++) {
+                const float * src = (const float *)((const char *) q->data + iq*q->nb[1] + h*q->nb[2]);
+                float       * row = q_host + (h*n_q + iq)*d_k;
+                for (int64_t i0 = 0; i0 < d_k; i0++) {
+                    row[i0] = src[i0] * scale;
+                }
+            }
+        }
+    });
+
+    ggml_zdnn_parallel_for(n_threads, n_head_kv*n_kv, [&](int64_t r_start, int64_t r_end) {
+        for (int64_t r = r_start; r < r_end; r++) {
+            const int64_t hk = r / n_kv;
+            const int64_t ik = r % n_kv;
+            memcpy(k_host + r*d_k, (const char *) k->data + ik*k->nb[1] + hk*k->nb[2], d_k*sizeof(ggml_fp16_t));
+            memcpy(v_host + r*d_v, (const char *) v->data + ik*v->nb[1] + hk*v->nb[2], d_v*sizeof(ggml_fp16_t));
+        }
+    });
+
+    zdnn_tensor_desc q_pre_tfm_desc, q_tfm_desc;
+    zdnn_tensor_desc k_pre_tfm_desc, k_tfm_desc;
+    zdnn_tensor_desc v_pre_tfm_desc, v_tfm_desc;
+    zdnn_tensor_desc o_pre_tfm_desc, o_tfm_desc;
+    zdnn_ztensor     q_ztensor, k_ztensor, v_ztensor, o_ztensor;
+
+    ggml_zdnn_pool_alloc<uint8_t> q_ztensor_alloc, k_ztensor_alloc, v_ztensor_alloc, o_ztensor_alloc;
+
+    // a single query keeps the heads of a kv head in one stack, which is the same memory order as one head per stack
+    if (single) {
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &q_pre_tfm_desc, n_head_kv, group, d_k);
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &o_pre_tfm_desc, n_head_kv, group, d_v);
+    } else {
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &q_pre_tfm_desc, n_head, n_q, d_k);
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &o_pre_tfm_desc, n_head, n_q, d_v);
+    }
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP16, &k_pre_tfm_desc, n_head_kv, n_kv, d_k);
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP16, &v_pre_tfm_desc, n_head_kv, n_kv, d_v);
+
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&q_pre_tfm_desc, &q_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&k_pre_tfm_desc, &k_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&v_pre_tfm_desc, &v_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&o_pre_tfm_desc, &o_tfm_desc));
+
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &q_pre_tfm_desc, &q_tfm_desc, &q_ztensor, q_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &k_pre_tfm_desc, &k_tfm_desc, &k_ztensor, k_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &v_pre_tfm_desc, &v_tfm_desc, &v_ztensor, v_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &o_pre_tfm_desc, &o_tfm_desc, &o_ztensor, o_ztensor_alloc);
+
+    ZDNN_CHECK(zdnn_transform_ztensor(&q_ztensor, q_host));
+    ZDNN_CHECK(zdnn_transform_ztensor(&k_ztensor, k_host));
+    ZDNN_CHECK(zdnn_transform_ztensor(&v_ztensor, v_host));
+
+    ggml_zdnn_pool_alloc<uint8_t> softmax_save_area(ctx->pool, ZDNN_SOFTMAX_SAVEAREA_SIZE);
+
+    zdnn_tensor_desc s_pre_tfm_desc, s_tfm_desc;
+    zdnn_ztensor     s_ztensor, s2_ztensor;
+
+    ggml_zdnn_pool_alloc<uint8_t> s_ztensor_alloc, s2_ztensor_alloc;
+
+    if (single) {
+        // the bias of K*Q is the mask row of the query, repeated for every kv head
+        zdnn_tensor_desc bias_pre_tfm_desc, bias_tfm_desc;
+        zdnn_ztensor     bias_ztensor;
+
+        ggml_zdnn_pool_alloc<uint8_t> bias_ztensor_alloc;
+
+        const zdnn_ztensor * bias = &ggml_zdnn_zero_bias(ctx, n_head_kv, n_kv);
+        if (mask != nullptr) {
+            const ggml_fp16_t * mask_row = (const ggml_fp16_t *) mask->data;
+            for (int64_t hk = 0; hk < n_head_kv; hk++) {
+                for (int64_t ik = 0; ik < n_kv; ik++) {
+                    mask_host[hk*n_kv + ik] = ggml_zdnn_fattn_mask_value(mask_row[ik]);
+                }
+            }
+
+            zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias_pre_tfm_desc, n_head_kv, n_kv);
+            ZDNN_CHECK(zdnn_generate_transformed_desc(&bias_pre_tfm_desc, &bias_tfm_desc));
+            ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &bias_pre_tfm_desc, &bias_tfm_desc, &bias_ztensor, bias_ztensor_alloc);
+            ZDNN_CHECK(zdnn_transform_ztensor(&bias_ztensor, mask_host));
+            bias = &bias_ztensor;
+        }
+
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &s_pre_tfm_desc, n_head_kv, group, n_kv);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&s_pre_tfm_desc, &s_tfm_desc));
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s_ztensor,  s_ztensor_alloc);
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s2_ztensor, s2_ztensor_alloc);
+
+        ZDNN_CHECK(zdnn_matmul_transpose_op(&q_ztensor, &k_ztensor, bias, false, true, MATMUL_OP_ADDITION, &s_ztensor));
+        ZDNN_CHECK(zdnn_softmax(&s_ztensor, softmax_save_area.get(), SOFTMAX_ACT_NONE, &s2_ztensor));
+        ZDNN_CHECK(zdnn_matmul_op(&s2_ztensor, &v_ztensor, &ggml_zdnn_zero_bias(ctx, n_head_kv, d_v), MATMUL_OP_ADDITION, &o_ztensor));
+    } else {
+        zdnn_tensor_desc mask_pre_tfm_desc, mask_tfm_desc;
+        zdnn_ztensor     mask_ztensor;
+
+        ggml_zdnn_pool_alloc<uint8_t> mask_ztensor_alloc;
+
+        if (mask != nullptr) {
+            ggml_zdnn_parallel_for(n_threads, n_q, [&](int64_t iq_start, int64_t iq_end) {
+                for (int64_t iq = iq_start; iq < iq_end; iq++) {
+                    const ggml_fp16_t * mask_row = (const ggml_fp16_t *)((const char *) mask->data + iq*mask->nb[1]);
+                    for (int64_t ik = 0; ik < n_kv; ik++) {
+                        mask_host[iq*n_kv + ik] = ggml_zdnn_fattn_mask_value(mask_row[ik]);
+                    }
+                }
+            });
+
+            zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &mask_pre_tfm_desc, 1, n_q, n_kv);
+            ZDNN_CHECK(zdnn_generate_transformed_desc(&mask_pre_tfm_desc, &mask_tfm_desc));
+            ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &mask_pre_tfm_desc, &mask_tfm_desc, &mask_ztensor, mask_ztensor_alloc);
+            ZDNN_CHECK(zdnn_transform_ztensor(&mask_ztensor, mask_host));
+        }
+
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &s_pre_tfm_desc, 1, n_q, n_kv);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&s_pre_tfm_desc, &s_tfm_desc));
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s_ztensor,  s_ztensor_alloc);
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s2_ztensor, s2_ztensor_alloc);
+
+        const zdnn_ztensor & k_bias = ggml_zdnn_zero_bias(ctx, 1, n_kv);
+        const zdnn_ztensor & v_bias = ggml_zdnn_zero_bias(ctx, 1, d_v);
+
+        zdnn_tensor_desc view_pre_tfm_desc[4];
+        zdnn_tensor_desc view_tfm_desc[4];
+
+        for (int64_t h = 0; h < n_head; h++) {
+            const zdnn_ztensor q_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[0], &view_tfm_desc[0], q_ztensor, n_head,    h,         1);
+            const zdnn_ztensor k_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[1], &view_tfm_desc[1], k_ztensor, n_head_kv, h / group, 1);
+            const zdnn_ztensor v_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[2], &view_tfm_desc[2], v_ztensor, n_head_kv, h / group, 1);
+            zdnn_ztensor       o_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[3], &view_tfm_desc[3], o_ztensor, n_head,    h,         1);
+
+            s_ztensor.is_transformed  = false;
+            s2_ztensor.is_transformed = false;
+
+            ZDNN_CHECK(zdnn_matmul_transpose_op(&q_view, &k_view, &k_bias, false, true, MATMUL_OP_ADDITION, &s_ztensor));
+
+            zdnn_ztensor * scores = &s_ztensor;
+            zdnn_ztensor * probs  = &s2_ztensor;
+            if (mask != nullptr) {
+                ZDNN_CHECK(zdnn_add(&s_ztensor, &mask_ztensor, &s2_ztensor));
+                scores = &s2_ztensor;
+                probs  = &s_ztensor;
+                probs->is_transformed = false;
+            }
+
+            ZDNN_CHECK(zdnn_softmax(scores, softmax_save_area.get(), SOFTMAX_ACT_NONE, probs));
+            ZDNN_CHECK(zdnn_matmul_op(probs, &v_view, &v_bias, MATMUL_OP_ADDITION, &o_view));
+        }
+
+        o_ztensor.is_transformed = true;
+    }
+
+    ZDNN_CHECK(zdnn_transform_origtensor(&o_ztensor, o_host));
+
+    // the result is [d_v][n_head][n_q]
+    for (int64_t iq = 0; iq < n_q; iq++) {
+        for (int64_t h = 0; h < n_head; h++) {
+            memcpy((char *) dst->data + h*dst->nb[1] + iq*dst->nb[2], o_host + (h*n_q + iq)*d_v, d_v*sizeof(float));
+        }
+    }
+
+    // the result is written to dst->data directly, so the dst ztensor is stale
+    ggml_backend_zdnn_buffer * dst_extra = (ggml_backend_zdnn_buffer *) dst->extra;
+    if (dst_extra != nullptr) {
+        dst_extra->ztensor.is_transformed = false;
+    }
+}

ggml/src/ggml-zdnn/fattn.hpp

@@ -0,0 +1,10 @@
+#ifndef GGML_ZDNN_FATTN_HPP
+#define GGML_ZDNN_FATTN_HPP
+
+#include "common.hpp"
+
+bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op);
+
+void ggml_zdnn_flash_attn_ext(ggml_backend_zdnn_context * ctx, ggml_tensor * dst);
+
+#endif  // GGML_ZDNN_FATTN_HPP

ggml/src/ggml-zdnn/ggml-zdnn.cpp

@@ -3,6 +3,7 @@
 #include "ggml-backend-impl.h"
 
 #include "ggml-zdnn/common.hpp"
+#include "ggml-zdnn/fattn.hpp"
 #include "ggml-zdnn/mmf.hpp"
 #include "ggml-zdnn/mmq.hpp"
 #include "ggml-zdnn/utils.hpp"
@@ -37,6 +38,11 @@ static bool ggml_zdnn_compute_forward(
                 ggml_zdnn_compute_forward_mul_mat(ctx, dst);
             } break;
 
+        case GGML_OP_FLASH_ATTN_EXT:
+            {
+                ggml_zdnn_flash_attn_ext(ctx, dst);
+            } break;
+
         default:
             return false;
     }
@@ -120,6 +126,9 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
                 }
             } break;
 
+        case GGML_OP_FLASH_ATTN_EXT:
+            return ggml_zdnn_supports_flash_attn_ext(ctx_dev, op);
+
         default:
             return false;
     }
@@ -205,9 +214,6 @@ static void ggml_zdnn_free(ggml_backend_zdnn_context * ctx) {
     for (const auto & it : ctx->zero_bias) {
         ZDNN_CHECK(zdnn_free_ztensor_buffer(&it.second->ztensor));
     }
-    for (void * scratch : ctx->scratch_ztensor) {
-        free(scratch);
-    }
 
     delete ctx;
 }
@@ -309,7 +315,19 @@ static void ggml_backend_zdnn_buffer_memset_tensor(ggml_backend_buffer_t buffer,
 static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     memcpy((char *)tensor->data + offset, data, size);
 
+    // views share the ztensor of their source, which is now stale, and empty tensors have none
+    if (tensor->view_src != nullptr) {
+        ggml_backend_zdnn_buffer * src_extra = (ggml_backend_zdnn_buffer *)tensor->view_src->extra;
+        if (src_extra != nullptr && src_extra->ztensor.is_transformed) {
+            zdnn_reset_ztensor(&src_extra->ztensor);
+        }
+        return;
+    }
+
     ggml_backend_zdnn_buffer * extra = (ggml_backend_zdnn_buffer *)tensor->extra;
+    if (extra == nullptr) {
+        return;
+    }
 
     if (ggml_is_quantized(tensor->type)) {
         GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor));
@@ -319,8 +337,8 @@ static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
 
     // Fixes the LLAMA_SET_ROWS bug
     // see: https://github.com/ggml-org/llama.cpp/issues/15414
-    // compute tensors are transformed on first use instead, as not every op reads the ztensor
-    if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+    // only weights are transformed here, others on first use, as not every op reads the ztensor and some data (masks) is not finite
+    if (tensor->buffer->usage != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
         zdnn_reset_ztensor(&extra->ztensor);
         return;
     }

ggml/src/ggml-zdnn/mmq.cpp

@@ -6,63 +6,6 @@
 #include <future>
 #include <vector>
 
-// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
-// the view descs must outlive the view
-static zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
-                                         zdnn_tensor_desc * tfm_desc,
-                                   const zdnn_ztensor     & src,
-                                         int64_t            n_stacks,
-                                         int64_t            s0,
-                                         int64_t            n) {
-
-    GGML_ASSERT(src.buffer_size % n_stacks == 0);
-    const uint64_t stack_bytes = src.buffer_size / n_stacks;
-
-    *pre_tfm_desc = *src.pre_transformed_desc;
-    *tfm_desc     = *src.transformed_desc;
-
-    switch (pre_tfm_desc->layout) {
-        case ZDNN_3DS:
-            {
-                pre_tfm_desc->dim3 = n;
-            } break;
-        case ZDNN_2DS:
-            {
-                pre_tfm_desc->dim2 = n;
-            } break;
-        default:
-            GGML_ABORT("%s: unsupported stacked layout", __func__);
-    }
-
-    tfm_desc->dim4 = n;
-    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
-
-    zdnn_ztensor view = src;
-    view.pre_transformed_desc = pre_tfm_desc;
-    view.transformed_desc     = tfm_desc;
-    view.buffer               = (char *) src.buffer + s0*stack_bytes;
-    view.buffer_size          = n * stack_bytes;
-    return view;
-}
-
-// zero bias of n_stacks stacks, created once per shape as it never changes
-static const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
-    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
-
-    if (!bias) {
-        bias = std::make_unique<ggml_backend_zdnn_buffer>();
-
-        std::vector<float> zeros(n_stacks * ne0, 0.0f);
-
-        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
-        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
-        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
-        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
-    }
-
-    return bias->ztensor;
-}
-
 // the weights are stacked into groups along ne00, and each group of each row of the inputs and weights has its own scale
 // zDNN multiplies the groups in stacked calls, then the partial results of each group are scaled and summed here
 static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
@@ -81,18 +24,20 @@ static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
     GGML_ASSERT(n_groups * gs >= ne10 && (n_groups - 1) * gs < ne10);
 
     // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
-    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * n_groups * gs * ne11);
+    ggml_zdnn_pool_alloc<int8_t> inputs_q_alloc(ctx->pool, 2 * n_groups * gs * ne11);
+    int8_t * inputs_q = inputs_q_alloc.get();
     std::vector<float> inputs_scales(n_groups * ne11);
     ggml_zdnn_quantize_inputs(n_threads, inputs, gs, weights_extra->channel_scales.data(),
                             inputs_q, inputs_q + n_groups * gs * ne11, inputs_scales.data());
 
-    zdnn_tensor_desc inputs_pre_tfm_desc;
-    zdnn_tensor_desc inputs_tfm_desc;
-    zdnn_ztensor     inputs_ztensor;
+    zdnn_tensor_desc              inputs_pre_tfm_desc;
+    zdnn_tensor_desc              inputs_tfm_desc;
+    zdnn_ztensor                  inputs_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> inputs_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2 * n_groups, ne11, gs);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, inputs_ztensor_alloc);
     ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
 
     // each group of the partial results is read back as its own 2D ztensor
@@ -116,19 +61,21 @@ static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
 
     // two partial buffers, so the next chunk runs on the NNPA while this one is summed
     // each holds the results of a chunk of inputs, followed by the results of their rounding errors
-    zdnn_tensor_desc partial_pre_tfm_desc;
-    zdnn_tensor_desc partial_tfm_desc;
-    zdnn_ztensor     partial_ztensor[2];
+    zdnn_tensor_desc              partial_pre_tfm_desc;
+    zdnn_tensor_desc              partial_tfm_desc;
+    zdnn_ztensor                  partial_ztensor[2];
+    ggml_zdnn_pool_alloc<uint8_t> partial_ztensor_alloc[2];
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2 * n_chunk, ne11, ne0);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
     GGML_ASSERT(zdnn_getsize_ztensor(&partial_tfm_desc) == (uint64_t)(2 * n_chunk) * group_bytes);
     for (int32_t i = 0; i < 2; i++) {
-        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], 1 + i);
+        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], partial_ztensor_alloc[i]);
     }
 
     const float * weights_scales = weights_extra->scales.data();
-    float * partial     = ggml_zdnn_resize_scratch_max(ctx->scratch_f32, 2 * ne11 * ne0);
+    ggml_zdnn_pool_alloc<float> partial_alloc(ctx->pool, 2 * ne11 * ne0);
+    float * partial     = partial_alloc.get();
     float * partial_res = partial + ne11 * ne0;
 
     const auto sum_chunk = [&](const zdnn_ztensor & chunk, int64_t g0, int64_t n) {
@@ -262,27 +209,30 @@ void ggml_zdnn_mul_mat_q(
 
     // zDNN takes one scale per zTensor, so quantize each row of the inputs with its own scale
     // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
-    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * ne10 * ne11);
+    ggml_zdnn_pool_alloc<int8_t> inputs_q_alloc(ctx->pool, 2 * ne10 * ne11);
+    int8_t * inputs_q = inputs_q_alloc.get();
     std::vector<float> inputs_scales(ne11);
     ggml_zdnn_quantize_inputs(n_threads, inputs, ne10, weights_extra->channel_scales.data(),
                             inputs_q, inputs_q + ne10 * ne11, inputs_scales.data());
 
-    zdnn_tensor_desc inputs_pre_tfm_desc;
-    zdnn_tensor_desc inputs_tfm_desc;
-    zdnn_ztensor     inputs_ztensor;
+    zdnn_tensor_desc              inputs_pre_tfm_desc;
+    zdnn_tensor_desc              inputs_tfm_desc;
+    zdnn_ztensor                  inputs_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> inputs_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2, ne11, ne10);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, inputs_ztensor_alloc);
     ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
 
-    zdnn_tensor_desc partial_pre_tfm_desc;
-    zdnn_tensor_desc partial_tfm_desc;
-    zdnn_ztensor     partial_ztensor;
+    zdnn_tensor_desc              partial_pre_tfm_desc;
+    zdnn_tensor_desc              partial_tfm_desc;
+    zdnn_ztensor                  partial_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> partial_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2, ne11, ne0);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, 1);
+    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, partial_ztensor_alloc);
 
     // both stacks are multiplied by the same unstacked weights in one broadcast call
     // the bias is all zeros, so it is already the pre-computed bias
@@ -300,7 +250,8 @@ void ggml_zdnn_mul_mat_q(
     const uint64_t stack_bytes = zdnn_getsize_ztensor(&stack_tfm_desc);
     GGML_ASSERT(partial_ztensor.buffer_size == 2*stack_bytes);
 
-    float * partial[2] = { (float *)output->data, ggml_zdnn_resize_scratch_max(ctx->scratch_f32, ne11 * ne0) };
+    ggml_zdnn_pool_alloc<float> partial_res_alloc(ctx->pool, ne11 * ne0);
+    float * partial[2] = { (float *)output->data, partial_res_alloc.get() };
 
     ggml_zdnn_parallel_for(n_threads, 2, [&](int64_t s_start, int64_t s_end) {
         zdnn_ztensor stack_ztensor;

ggml/src/ggml-zdnn/scratch.cpp

@@ -0,0 +1,45 @@
+#include "ggml.h"
+#include "scratch.hpp"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+
+ggml_zdnn_pool::~ggml_zdnn_pool() {
+    for (const buffer & b : cached) {
+        std::free(b.ptr);
+    }
+}
+
+void * ggml_zdnn_pool::alloc(size_t size, size_t * actual_size) {
+    int64_t ibest = -1;
+    for (int64_t i = 0; i < (int64_t)cached.size(); i++) {
+        if (cached[i].size >= size && (ibest < 0 || cached[i].size < cached[ibest].size)) {
+            ibest = i;
+        }
+    }
+
+    if (ibest >= 0) {
+        void * ptr   = cached[ibest].ptr;
+        *actual_size = cached[ibest].size;
+        cached.erase(cached.begin() + ibest);
+        return ptr;
+    }
+
+    // every cached buffer is too small, free the largest so the pool does not grow while sizes grow with the context
+    if (!cached.empty()) {
+        const auto largest = std::max_element(cached.begin(), cached.end(), [](const buffer & a, const buffer & b) { return a.size < b.size; });
+        std::free(largest->ptr);
+        cached.erase(largest);
+    }
+
+    // 5% more, so slowly growing sizes do not allocate every time
+    *actual_size = GGML_PAD(std::max<size_t>(1, (size_t)(1.05 * size)), 4096);
+    void * ptr   = std::aligned_alloc(4096, *actual_size);
+    GGML_ASSERT(ptr != nullptr);
+    return ptr;
+}
+
+void ggml_zdnn_pool::free(void * ptr, size_t size) {
+    cached.push_back({ ptr, size });
+}

ggml/src/ggml-zdnn/scratch.hpp

@@ -0,0 +1,61 @@
+#ifndef GGML_ZDNN_SCRATCH_HPP
+#define GGML_ZDNN_SCRATCH_HPP
+
+#include "ggml.h"
+
+#include <cstddef>
+#include <vector>
+
+// 4K aligned host buffers that are cached for reuse, see ggml_cuda_pool_leg
+struct ggml_zdnn_pool {
+    struct buffer {
+        void * ptr  = nullptr;
+        size_t size = 0;
+    };
+
+    std::vector<buffer> cached;
+
+    ~ggml_zdnn_pool();
+
+    void * alloc(size_t size, size_t * actual_size);
+    void   free(void * ptr, size_t size);
+};
+
+// a buffer of the pool that goes back to it at the end of its scope, see ggml_cuda_pool_alloc
+template <typename T>
+struct ggml_zdnn_pool_alloc {
+    ggml_zdnn_pool * pool = nullptr;
+    T * ptr = nullptr;
+    size_t actual_size = 0;
+
+    ggml_zdnn_pool_alloc() = default;
+
+    ggml_zdnn_pool_alloc(ggml_zdnn_pool & pool, size_t size) {
+        alloc(pool, size);
+    }
+
+    ~ggml_zdnn_pool_alloc() {
+        if (ptr != nullptr) {
+            pool->free(ptr, actual_size);
+        }
+    }
+
+    // size is in number of elements
+    T * alloc(ggml_zdnn_pool & pool, size_t size) {
+        GGML_ASSERT(ptr == nullptr);
+        this->pool = &pool;
+        ptr = (T *) pool.alloc(size * sizeof(T), &actual_size);
+        return ptr;
+    }
+
+    T * get() {
+        return ptr;
+    }
+
+    ggml_zdnn_pool_alloc(const ggml_zdnn_pool_alloc &) = delete;
+    ggml_zdnn_pool_alloc(ggml_zdnn_pool_alloc &&) = delete;
+    ggml_zdnn_pool_alloc & operator=(const ggml_zdnn_pool_alloc &) = delete;
+    ggml_zdnn_pool_alloc & operator=(ggml_zdnn_pool_alloc &&) = delete;
+};
+
+#endif  // GGML_ZDNN_SCRATCH_HPP