From 0a1dd13c95d8f4f1554c6c32fa21870ad95c050d Mon Sep 17 00:00:00 2001
From: cora4 <80816140+cora4@users.noreply.github.com>
Date: Sun, 12 Jul 2026 06:26:38 +0100
Subject: [PATCH 1/4] Set GGML_AVXVNNI to OFF by default (#2116)

---
 ggml/CMakeLists.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt
index f42e4569b..ec87b0807 100644
--- a/ggml/CMakeLists.txt
+++ b/ggml/CMakeLists.txt
@@ -82,7 +82,7 @@ option(GGML_CPU_HBM     "ggml: use memkind for CPU HBM" OFF)
 
 option(GGML_AVX         "ggml: enable AVX"              ${INS_ENB})
 option(GGML_AVX2        "ggml: enable AVX2"             ${INS_ENB})
-option(GGML_AVXVNNI     "ggml: enable AVX-VNNI"         ${INS_ENB})
+option(GGML_AVXVNNI     "ggml: enable AVX-VNNI"         OFF)
 option(GGML_AVX512      "ggml: enable AVX512"           OFF)
 option(GGML_AVX512_VBMI "ggml: enable AVX512-VBMI"      OFF)
 option(GGML_AVX512_VNNI "ggml: enable AVX512-VNNI"      OFF)

From 97bc869552d577818c847dee27c017d661054ec1 Mon Sep 17 00:00:00 2001
From: Joel Farthing <farthing@me.com>
Date: Sun, 12 Jul 2026 01:59:50 -0500
Subject: [PATCH 2/4] ggml: add fused sinkhorn op (eps + output-layout params)
 (#2115)

* ggml: add fused sinkhorn op (eps + output-layout params); use it for openPangu mHC

* openpangu: call ggml_sinkhorn directly from mhc_post

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
---
 ggml/include/ggml.h             |  17 +++++
 ggml/src/ggml-cuda.cu           |   9 +++
 ggml/src/ggml-cuda/sinkhorn.cu  | 103 ++++++++++++++++++++++++++++
 ggml/src/ggml-cuda/sinkhorn.cuh |   3 +
 ggml/src/ggml.c                 | 115 +++++++++++++++++++++++++++++++-
 src/graphs/build_openpangu.cpp  |  38 +----------
 6 files changed, 247 insertions(+), 38 deletions(-)
 create mode 100644 ggml/src/ggml-cuda/sinkhorn.cu
 create mode 100644 ggml/src/ggml-cuda/sinkhorn.cuh

diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index 0f0d08ec2..40b2c7ca3 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -705,6 +705,7 @@ extern "C" {
         GGML_OP_FUSED_RMS_RMS_ADD,
         GGML_OP_BLEND,
         GGML_OP_INDEXER_TOPK,
+        GGML_OP_SINKHORN,
 
         GGML_OP_COUNT,
     };
@@ -2578,6 +2579,22 @@ extern "C" {
             enum ggml_unary_op    op,
             int                   n_top_k);
 
+    // Sinkhorn normalization of a flat [S*S, T] batch of S x S matrices into
+    // doubly-stochastic form: softmax over columns, then column normalization,
+    // then (n_iters - 1) rounds of row + column normalization (ends on columns).
+    // The flat input is row-major (column index fastest). eps, when non-zero, is
+    // added to the softmax output and to every normalization sum before dividing.
+    // With output_transposed the result is [S, S, T] with ne0 = row, ne1 = column
+    // (ready for out[c] = sum_r m[r,c] * residual[r] consumers); otherwise the
+    // bare input layout (ne0 = column) is kept.
+    GGML_API struct ggml_tensor * ggml_sinkhorn(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   S,
+            int                   n_iters,
+            float                 eps,
+            bool                  output_transposed);
+
     // custom operators
 
     typedef void (*ggml_unary_op_f32_t) (const int, float *, const float *);
diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu
index b9a0d2f26..cd88309ea 100644
--- a/ggml/src/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda.cu
@@ -56,6 +56,7 @@
 #include "ggml-cuda/reduce.cuh"
 #include "ggml-cuda/tri.cuh"
 #include "ggml-cuda/delta-net.cuh"
+#include "ggml-cuda/sinkhorn.cuh"
 #include "ggml-cuda/blend.cuh"
 #include "ggml-cuda/indexer_topk.cuh"
 
@@ -4127,6 +4128,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_DELTA_NET:
             ggml_cuda_op_delta_net(ctx, dst);
             break;
+        case GGML_OP_SINKHORN:
+            ggml_cuda_op_sinkhorn(ctx, dst);
+            break;
         case GGML_OP_FLASH_ATTN_EXT:
             ggml_cuda_flash_attn_ext(ctx, dst);
             break;
@@ -5038,6 +5042,11 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
         case GGML_OP_DELTA_NET:
         case GGML_OP_INDEXER_TOPK:
             return true;
+        case GGML_OP_SINKHORN: {
+            const int sink_s = op->op_params[0];
+            return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 &&
+                   sink_s >= 1 && sink_s <= 8 && op->src[0]->ne[0] == (int64_t) sink_s*sink_s;
+        }
         case GGML_OP_FLASH_ATTN_EXT:
 #if defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)
             return (op->src[0]->ne[0] == 64 && op->src[1]->type == GGML_TYPE_F16) || op->src[0]->ne[0] == 128;
diff --git a/ggml/src/ggml-cuda/sinkhorn.cu b/ggml/src/ggml-cuda/sinkhorn.cu
new file mode 100644
index 000000000..f02a8a440
--- /dev/null
+++ b/ggml/src/ggml-cuda/sinkhorn.cu
@@ -0,0 +1,103 @@
+#include "common.cuh"
+#include "sinkhorn.cuh"
+
+// Sinkhorn normalization of T independent S x S matrices (S <= 8, so a matrix is
+// at most 64 floats). One thread per token: the whole matrix lives in a thread-local
+// array and the 6-node-per-iteration graph chain collapses into this single kernel.
+// Semantics match the reference: softmax over columns, column normalization,
+// then (iters - 1) rounds of row + column normalization (ends on columns).
+// Input is the flat [S*S, T] row-major tensor (column index fastest); output is
+// [S, S, T] with ne0 = row, i.e. transposed on write.
+
+template <int S>
+static __global__ void k_sinkhorn(const float * __restrict__ x, float * __restrict__ dst,
+                                  const int64_t T, const int iters, const float eps,
+                                  const int transposed, const int64_t nb1) {
+    const int64_t t = (int64_t) blockIdx.x*blockDim.x + threadIdx.x;
+    if (t >= T) {
+        return;
+    }
+
+    const float * xt = (const float *)((const char *) x + t*nb1);
+    float m[S*S];
+
+    #pragma unroll
+    for (int r = 0; r < S; ++r) {
+        float mx = xt[r*S];
+        for (int c = 1; c < S; ++c) mx = fmaxf(mx, xt[r*S + c]);
+        float sum = 0.0f;
+        for (int c = 0; c < S; ++c) { m[r*S + c] = expf(xt[r*S + c] - mx); sum += m[r*S + c]; }
+        for (int c = 0; c < S; ++c) m[r*S + c] = m[r*S + c]/sum + eps;
+    }
+    #pragma unroll
+    for (int c = 0; c < S; ++c) {
+        float sum = eps;
+        for (int r = 0; r < S; ++r) sum += m[r*S + c];
+        for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+    }
+    for (int i = 0; i < iters - 1; ++i) {
+        #pragma unroll
+        for (int r = 0; r < S; ++r) {
+            float sum = eps;
+            for (int c = 0; c < S; ++c) sum += m[r*S + c];
+            for (int c = 0; c < S; ++c) m[r*S + c] /= sum;
+        }
+        #pragma unroll
+        for (int c = 0; c < S; ++c) {
+            float sum = eps;
+            for (int r = 0; r < S; ++r) sum += m[r*S + c];
+            for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+        }
+    }
+
+    float * yt = dst + t*S*S;
+    if (transposed) {
+        #pragma unroll
+        for (int c = 0; c < S; ++c) {
+            for (int r = 0; r < S; ++r) yt[c*S + r] = m[r*S + c];
+        }
+    } else {
+        #pragma unroll
+        for (int k = 0; k < S*S; ++k) yt[k] = m[k];
+    }
+}
+
+void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    const int S     = dst->op_params[0];
+    const int iters = dst->op_params[1];
+    float eps;
+    memcpy(&eps, &dst->op_params[2], sizeof(float));
+    const int transposed = dst->op_params[3];
+    const int64_t T = src0->ne[1];
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(S >= 1 && S <= 8);
+    GGML_ASSERT(src0->ne[0] == (int64_t) S * S);
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    if (T == 0) {
+        return;
+    }
+
+    const int block = 256;
+    const int64_t grid = (T + block - 1)/block;
+    cudaStream_t stream = ctx.stream();
+
+    const float * x = (const float *) src0->data;
+    float * y = (float *) dst->data;
+
+    switch (S) {
+        case 1: k_sinkhorn<1><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 2: k_sinkhorn<2><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 3: k_sinkhorn<3><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 4: k_sinkhorn<4><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 5: k_sinkhorn<5><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 6: k_sinkhorn<6><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 7: k_sinkhorn<7><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 8: k_sinkhorn<8><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        default: GGML_ABORT("sinkhorn: unsupported S");
+    }
+}
diff --git a/ggml/src/ggml-cuda/sinkhorn.cuh b/ggml/src/ggml-cuda/sinkhorn.cuh
new file mode 100644
index 000000000..c9f86186c
--- /dev/null
+++ b/ggml/src/ggml-cuda/sinkhorn.cuh
@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index 79f89b575..3fe4e6d43 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -4333,9 +4333,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "FUSED_RMS_RMS_ADD",
     "BLEND",
     "INDEXER_TOPK",
+    "SINKHORN",
 };
 
-static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104");
+static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -4455,10 +4456,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "rms(x1)+rms(x2)",
     "blend(a,b,c)",
     "indexer_topk(k, q, w, mask)",
+    "sinkhorn(x)",
 
 };
 
-static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104");
+static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -10139,6 +10141,32 @@ struct ggml_tensor * ggml_indexer_topk(
 }
 
 
+struct ggml_tensor * ggml_sinkhorn(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   S,
+            int                   n_iters,
+            float                 eps,
+            bool                  output_transposed) {
+    GGML_ASSERT(eps >= 0.0f);
+    GGML_ASSERT(a->type == GGML_TYPE_F32);
+    GGML_ASSERT(S >= 1 && S <= 8);
+    GGML_ASSERT(a->ne[0] == (int64_t) S * S);
+    GGML_ASSERT(a->ne[2] == 1 && a->ne[3] == 1);
+    GGML_ASSERT(n_iters >= 1);
+    GGML_ASSERT(ggml_is_contiguous(a));
+
+    struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, S, a->ne[1]);
+    result->op = GGML_OP_SINKHORN;
+    result->op_params[0] = S;
+    result->op_params[1] = n_iters;
+    memcpy(&result->op_params[2], &eps, sizeof(float));
+    result->op_params[3] = output_transposed ? 1 : 0;
+    result->src[0] = a;
+
+    return result;
+}
+
 // ggml_fill
 
 static struct ggml_tensor * ggml_fill_impl(
@@ -23086,6 +23114,83 @@ static void ggml_compute_forward_delta_net(
     }
 }
 
+// ggml_compute_forward_sinkhorn
+
+static void ggml_compute_forward_sinkhorn_f32(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0];
+
+    const int S     = dst->op_params[0];
+    const int iters = dst->op_params[1];
+    float eps;
+    memcpy(&eps, &dst->op_params[2], sizeof(float));
+    const int transposed = dst->op_params[3];
+    const int64_t T = src0->ne[1];
+
+    GGML_ASSERT(S >= 1 && S <= 8);
+    GGML_ASSERT(src0->ne[0] == (int64_t) S * S);
+    GGML_ASSERT(iters >= 1);
+
+    // one token is S*S floats (16 at S=4): parallelize over tokens only
+    const int64_t t0 = (T *  params->ith   ) / params->nth;
+    const int64_t t1 = (T * (params->ith+1)) / params->nth;
+
+    float m[64];
+
+    for (int64_t t = t0; t < t1; ++t) {
+        const float * x = (const float *)((const char *)src0->data + t*src0->nb[1]);
+        float       * y = (float       *)((      char *)dst->data  + t*dst->nb[2]);
+
+        // softmax over columns c for each row r; flat input is row-major (c fastest)
+        for (int r = 0; r < S; ++r) {
+            float mx = x[r*S];
+            for (int c = 1; c < S; ++c) mx = MAX(mx, x[r*S + c]);
+            float sum = 0.0f;
+            for (int c = 0; c < S; ++c) { m[r*S + c] = expf(x[r*S + c] - mx); sum += m[r*S + c]; }
+            for (int c = 0; c < S; ++c) m[r*S + c] = m[r*S + c]/sum + eps;
+        }
+        // column normalization first, then (iters - 1) rounds of row + column: ends on columns
+        for (int c = 0; c < S; ++c) {
+            float sum = eps;
+            for (int r = 0; r < S; ++r) sum += m[r*S + c];
+            for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+        }
+        for (int i = 0; i < iters - 1; ++i) {
+            for (int r = 0; r < S; ++r) {
+                float sum = eps;
+                for (int c = 0; c < S; ++c) sum += m[r*S + c];
+                for (int c = 0; c < S; ++c) m[r*S + c] /= sum;
+            }
+            for (int c = 0; c < S; ++c) {
+                float sum = eps;
+                for (int r = 0; r < S; ++r) sum += m[r*S + c];
+                for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+            }
+        }
+        if (transposed) {
+            // dst is [row, col, T] (ne0 = row): transpose on write
+            for (int c = 0; c < S; ++c) {
+                for (int r = 0; r < S; ++r) y[c*S + r] = m[r*S + c];
+            }
+        } else {
+            for (int k = 0; k < S*S; ++k) y[k] = m[k];
+        }
+    }
+}
+
+static void ggml_compute_forward_sinkhorn(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    switch (dst->src[0]->type) {
+        case GGML_TYPE_F32:
+            ggml_compute_forward_sinkhorn_f32(params, dst);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
 // ggml_compute_forward_win_part
 
 static void ggml_compute_forward_win_part_f32(
@@ -24828,6 +24933,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml
             {
                 ggml_compute_forward_delta_net(params, tensor);
             } break;
+        case GGML_OP_SINKHORN:
+            {
+                ggml_compute_forward_sinkhorn(params, tensor);
+            } break;
         case GGML_OP_INDEXER_TOPK:
             {
                 if (!iqk_indexer_topk(tensor, params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, params->ith, params->nth)) {
@@ -25896,6 +26005,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_DELTA_NET:
         case GGML_OP_INDEXER_TOPK:
+        case GGML_OP_SINKHORN:
             {
                 GGML_ABORT("fatal error"); // TODO: not implemented
             }
@@ -26633,6 +26743,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_DELTA_NET:
         case GGML_OP_INDEXER_TOPK:
+        case GGML_OP_SINKHORN:
             {
                 n_tasks = n_threads;
             } break;
diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp
index 4dfb1bdd2..37371f137 100644
--- a/src/graphs/build_openpangu.cpp
+++ b/src/graphs/build_openpangu.cpp
@@ -222,39 +222,6 @@ static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf,
     return out;
 }
 
-// --- mHC Sinkhorn: h_res [S*S, T] -> doubly-stochastic per token, 20 iters (ends on col norm) ---
-static ggml_tensor * openpangu_sinkhorn(ggml_context * ctx, ggml_tensor * h_res_flat,
-                                        int64_t S, int64_t T, int iters, float hc_eps) {
-    // The flat h_res is torch [r,c] row-major (c fastest), so a bare reshape gives ne0=col.
-    // Transpose once so ne0=row(S), ne1=col(S): every axis op below then matches the
-    // reference _mhc_sinkhorn_naive (softmax over col, first norm over row, end on col-sum=1)
-    // and mhc_post's out[c] = sum_r m[r,c]*residual[r].
-    (void) hc_eps; // softmax outputs are strictly positive, so the eps is numerically inert here
-    ggml_tensor * m = ggml_reshape_3d(ctx, h_res_flat, S, S, T);   // ne0=col (bare reshape)
-    // ref softmaxes h_res over columns; a bare reshape already has ne0=col, so soft_max
-    // (over ne0) hits the column axis directly -- no pre-permute round-trip needed.
-    m = ggml_soft_max(ctx, m);                                 // softmax over col
-    m = ggml_cont(ctx, ggml_permute(ctx, m, 1, 0, 2, 3));      // transpose once -> [row, col, T]
-
-    auto col_norm = [&](ggml_tensor * a) {
-        ggml_tensor * col_sum = ggml_sum_rows(ctx, a);          // sums ne0(row) -> [1, col, T]
-        return ggml_div(ctx, a, col_sum);                       // broadcast [1,col,T] over rows
-    };
-    auto row_norm = [&](ggml_tensor * a) {
-        ggml_tensor * ap = ggml_cont(ctx, ggml_permute(ctx, a, 1, 0, 2, 3)); // [col,row,T]
-        ggml_tensor * row_sum = ggml_sum_rows(ctx, ap);         // [1, row, T]
-        ggml_tensor * out = ggml_div(ctx, ap, row_sum);
-        return ggml_cont(ctx, ggml_permute(ctx, out, 1, 0, 2, 3)); // back [row,col,T]
-    };
-
-    m = col_norm(m);
-    for (int i = 0; i < iters - 1; ++i) {
-        m = row_norm(m);
-        m = col_norm(m);
-    }
-    return m; // [row(S), col(S), T]
-}
-
 // Attention sublayer body, shared by the base layers and the NextN/MTP head.
 // x_normed = input-layernormed hidden [n_embd, T]; returns post-o_proj output [n_embd, T].
 // conv_state is the recurrent MoME state slot. seq_qnext is the [1, T] sequence-id input
@@ -876,7 +843,6 @@ ggml_cgraph * llm_build_context::build_openpangu() {
     const int64_t n_embd_head_k = hparams.n_embd_head_k(0);                // 192
     const int64_t S             = hparams.mhc_num_stream;                  // 4
     const int    sink_iters     = (int) hparams.mhc_recur_norm;            // 20
-    const float  hc_eps         = 1e-6f;
     const float  kq_scale       = 1.0f / sqrtf(float(n_embd_head_k));
 
 
@@ -1076,7 +1042,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
         // cont is required: the CUDA broadcast-mul path misreads strided views (h_pre is a
         // row-slice of mixes), while CPU handles the strides — token 0 right, tokens 1+ garbage
         h_pre = ggml_add(ctx0, ggml_mul(ctx0, ggml_cont(ctx0, h_pre), a_pre), b_pre);  // broadcast scalar + [S]
-        h_pre = ggml_sigmoid(ctx0, h_pre);                            // [S,T] (+hc_eps omitted, inert)
+        h_pre = ggml_sigmoid(ctx0, h_pre);                            // [S,T] (+eps omitted, inert)
 
         // combine: x[h,t] = sum_s h_pre[s,t] * R[h,s,t]
         ggml_tensor * hpre3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_pre), 1, S, n_tokens);
@@ -1101,7 +1067,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
         h_post = ggml_scale(ctx0, ggml_sigmoid(ctx0, h_post), 2.0f);  // 2*sigmoid, [S,T]
 
         ggml_tensor * m = ggml_add(ctx0, ggml_mul(ctx0, h_res, a_res), b_res); // [S*S,T]
-        m = openpangu_sinkhorn(ctx0, m, S, n_tokens, sink_iters, hc_eps);      // [row S, col S, T]
+        m = ggml_sinkhorn(ctx0, m, (int) S, sink_iters, 0.0f, /*output_transposed=*/true); // [row S, col S, T]
 
         // term1: h_post[s,t]*y[h,t] -> [H,S,T]
         ggml_tensor * y3 = ggml_reshape_3d(ctx0, y, n_embd, 1, n_tokens);

From 8e2d83cfee5c2de18777d01a52ff9b73c34cd335 Mon Sep 17 00:00:00 2001
From: cora4 <80816140+cora4@users.noreply.github.com>
Date: Sun, 12 Jul 2026 09:58:09 +0100
Subject: [PATCH 3/4] Fix per_layer_token_embedding (#2117)

Co-authored-by: cora4 <you@example.com>
---
 examples/quantize/quantize.cpp |  6 ++++
 ggml/src/iqk/iqk_quantize.cpp  |  5 +--
 include/llama.h                |  1 +
 src/llama-quantize.cpp         | 35 +++++++++++++++++++++
 src/llama.cpp                  | 57 +++++++++++++++++-----------------
 5 files changed, 72 insertions(+), 32 deletions(-)

diff --git a/examples/quantize/quantize.cpp b/examples/quantize/quantize.cpp
index 3734a67d6..621b5adf5 100644
--- a/examples/quantize/quantize.cpp
+++ b/examples/quantize/quantize.cpp
@@ -396,6 +396,12 @@ int main(int argc, char ** argv) {
             } else {
                 usage(argv[0]);
             }
+        } else if (strcmp(argv[arg_idx], "--per-layer-token-embedding-type") == 0) {
+            if (arg_idx < argc-1) {
+                params.per_layer_token_embedding_type = parse_ggml_type(argv[++arg_idx]);
+            } else {
+                usage(argv[0]);
+            }
         } else if (strcmp(argv[arg_idx], "--ffn-gate-inp-type") == 0) {
             if (arg_idx < argc-1) {
                 params.ffn_gate_inp_type = parse_ggml_type(argv[++arg_idx]);
diff --git a/ggml/src/iqk/iqk_quantize.cpp b/ggml/src/iqk/iqk_quantize.cpp
index b049aa584..35b21c2e2 100644
--- a/ggml/src/iqk/iqk_quantize.cpp
+++ b/ggml/src/iqk/iqk_quantize.cpp
@@ -8353,10 +8353,7 @@ const Modify * get_modify_info(ggml_type type) {
     return it != k_mod_map.end() ? &it->second : nullptr;
 }
 bool is_forbidden_tensor(const std::string& name) {
-    static const std::string kTokenEmbd{"token_embd.weight"};
-    if (name == kTokenEmbd) return true;
-    //if (auto pos = name.find("attn_kv_b.weight"); pos != std::string::npos) return true;
-    return false;
+    return (name == "token_embd.weight" || name == "per_layer_token_embd.weight");
 }
 }
 
diff --git a/include/llama.h b/include/llama.h
index ce5458335..05eb3b8fd 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -523,6 +523,7 @@ extern "C" {
         enum llama_ftype ftype;              // quantize to this llama_ftype
         enum ggml_type output_tensor_type;   // output tensor type
         enum ggml_type token_embedding_type; // token embeddings tensor type
+        enum ggml_type per_layer_token_embedding_type; // token embeddings tensor type
         enum ggml_type attn_q_type;          // attention query tensor type
         enum ggml_type attn_k_type;          // attention key tensor type
         enum ggml_type attn_v_type;          // attention value tensor type
diff --git a/src/llama-quantize.cpp b/src/llama-quantize.cpp
index 574ab0ce8..0c12e40e6 100644
--- a/src/llama-quantize.cpp
+++ b/src/llama-quantize.cpp
@@ -409,6 +409,29 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
                 new_type = GGML_TYPE_IQ4_NL;
             }
         }
+    } else if (name == "per_layer_token_embd.weight") {
+        if (qs.params->per_layer_token_embedding_type < GGML_TYPE_COUNT) {
+            new_type = qs.params->per_layer_token_embedding_type;
+        } else {
+            if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ2_XS ||
+                ftype == LLAMA_FTYPE_MOSTLY_IQ1_S   || ftype == LLAMA_FTYPE_MOSTLY_IQ1_M  ||
+                ftype == LLAMA_FTYPE_MOSTLY_IQ2_XXS_R4 || ftype == LLAMA_FTYPE_MOSTLY_IQ2_XS_R4 ||
+                ftype == LLAMA_FTYPE_MOSTLY_IQ1_S_R4 || ftype == LLAMA_FTYPE_MOSTLY_IQ1_M_R4) {
+                new_type = GGML_TYPE_Q2_K;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M_R4) {
+                new_type = GGML_TYPE_IQ3_S;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ3_KT) {
+                new_type = GGML_TYPE_IQ3_S;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS_R4) {
+                new_type = GGML_TYPE_IQ3_K;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ1_BN || ftype == LLAMA_FTYPE_MOSTLY_IQ2_BN || ftype == LLAMA_FTYPE_MOSTLY_IQ2_BN_R4) {
+                new_type = GGML_TYPE_IQ4_NL;
+            }
+        }
     } else if (ftype == LLAMA_FTYPE_MOSTLY_IQ1_S_R4 || ftype == LLAMA_FTYPE_MOSTLY_IQ1_M_R4) {
         if (name.find("attn_v.weight") != std::string::npos) {
             if (qs.model.hparams.n_expert >= 4 || qs.model.hparams.n_gqa() >= 4) new_type = GGML_TYPE_IQ4_K_R4;
@@ -830,6 +853,15 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
         }
     }
 
+    if (name == "per_layer_token_embd.weight") {
+        auto working_type = interleaved_properties(new_type).first;
+        if (working_type != new_type) {
+            printf("\n============ Per-layer Token embeddings cannot be quantized with row-interleaved quants\n");
+            printf("---> Changed %s to %s\n", ggml_type_name(new_type), ggml_type_name(working_type));
+            new_type = working_type;
+        }
+    }
+
     return new_type;
 }
 
@@ -1522,6 +1554,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
             if (params->token_embedding_type < GGML_TYPE_COUNT && strcmp(tensor->name, "token_embd.weight") == 0) {
                 new_type = params->token_embedding_type;
             }
+            if (params->per_layer_token_embedding_type < GGML_TYPE_COUNT && strcmp(tensor->name, "per_layer_token_embd.weight") == 0) {
+                new_type = params->per_layer_token_embedding_type;
+            }
             if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) {
                 new_type = params->output_tensor_type;
             }
diff --git a/src/llama.cpp b/src/llama.cpp
index 56c6b2bad..4987911c8 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -7186,34 +7186,35 @@ struct llama_context_params llama_context_default_params() {
 
 struct llama_model_quantize_params llama_model_quantize_default_params() {
     struct llama_model_quantize_params result = {
-        /*.nthread                     =*/ 0,
-        /*.ftype                       =*/ LLAMA_FTYPE_MOSTLY_Q5_1,
-        /*.output_tensor_type          =*/ GGML_TYPE_COUNT,
-        /*.token_embedding_type        =*/ GGML_TYPE_COUNT,
-        /*.attn_q_type                 =*/ GGML_TYPE_COUNT,
-        /*.attn_k_type                 =*/ GGML_TYPE_COUNT,
-        /*.attn_v_type                 =*/ GGML_TYPE_COUNT,
-        /*.attn_qkv_type               =*/ GGML_TYPE_COUNT,
-        /*.attn_output_type            =*/ GGML_TYPE_COUNT,
-        /*.ffn_gate_type               =*/ GGML_TYPE_COUNT,
-        /*.ffn_down_type               =*/ GGML_TYPE_COUNT,
-        /*.ffn_up_type                 =*/ GGML_TYPE_COUNT,
-        /*.ffn_gat_inp_type            =*/ GGML_TYPE_COUNT,
-        /*.extra_output_type           =*/ GGML_TYPE_COUNT,
-        /*.allow_requantize            =*/ false,
-        /*.quantize_output_tensor      =*/ true,
-        /*.only_copy                   =*/ false,
-        /*.pure                        =*/ false,
-        /*.keep_split                  =*/ false,
-        /*.ignore_imatrix_rules        =*/ false,
-        /*.only_repack                 =*/ false,
-        /*.dry_run                     =*/ false,
-        /*.partial_requant             =*/ false,
-        /*.imatrix                     =*/ nullptr,
-        /*.kv_overrides                =*/ nullptr,
-        /*.custom_quants               =*/ nullptr,
-        /*.repack_pattern              =*/ nullptr,
-        /*.user_data                   =*/ nullptr,
+        /*.nthread                        =*/ 0,
+        /*.ftype                          =*/ LLAMA_FTYPE_MOSTLY_Q5_1,
+        /*.output_tensor_type             =*/ GGML_TYPE_COUNT,
+        /*.token_embedding_type           =*/ GGML_TYPE_COUNT,
+        /*.per_layer_token_embedding_type =*/ GGML_TYPE_COUNT,
+        /*.attn_q_type                    =*/ GGML_TYPE_COUNT,
+        /*.attn_k_type                    =*/ GGML_TYPE_COUNT,
+        /*.attn_v_type                    =*/ GGML_TYPE_COUNT,
+        /*.attn_qkv_type                  =*/ GGML_TYPE_COUNT,
+        /*.attn_output_type               =*/ GGML_TYPE_COUNT,
+        /*.ffn_gate_type                  =*/ GGML_TYPE_COUNT,
+        /*.ffn_down_type                  =*/ GGML_TYPE_COUNT,
+        /*.ffn_up_type                    =*/ GGML_TYPE_COUNT,
+        /*.ffn_gat_inp_type               =*/ GGML_TYPE_COUNT,
+        /*.extra_output_type              =*/ GGML_TYPE_COUNT,
+        /*.allow_requantize               =*/ false,
+        /*.quantize_output_tensor         =*/ true,
+        /*.only_copy                      =*/ false,
+        /*.pure                           =*/ false,
+        /*.keep_split                     =*/ false,
+        /*.ignore_imatrix_rules           =*/ false,
+        /*.only_repack                    =*/ false,
+        /*.dry_run                        =*/ false,
+        /*.partial_requant                =*/ false,
+        /*.imatrix                        =*/ nullptr,
+        /*.kv_overrides                   =*/ nullptr,
+        /*.custom_quants                  =*/ nullptr,
+        /*.repack_pattern                 =*/ nullptr,
+        /*.user_data                      =*/ nullptr,
     };
 
     return result;

From 0d5d1330e2780f5f18b999e526cb3441275b97c1 Mon Sep 17 00:00:00 2001
From: Kawrakow <iwankawrakow@gmail.com>
Date: Sun, 12 Jul 2026 19:33:33 +0300
Subject: [PATCH 4/4] GLM-DSA: much better PP long context performance (CUDA)
 (#2109)

* WIP: indexer_topk on CUDA

* Forgot these

* WIP

* WIP

* This seems to work

* Minor

* Fix bug. Fix suggested by @sayap using GLM-5.2

* GLM-DSA: much better PP long context performance (CUDA)
---
 ggml/src/ggml-cuda/dsa_attn.cu  | 312 ++++++++++++++++++++++++++++++++
 ggml/src/ggml-cuda/dsa_attn.cuh |   3 +
 ggml/src/ggml-cuda/fattn.cu     |   8 +-
 3 files changed, 322 insertions(+), 1 deletion(-)
 create mode 100644 ggml/src/ggml-cuda/dsa_attn.cu
 create mode 100644 ggml/src/ggml-cuda/dsa_attn.cuh

diff --git a/ggml/src/ggml-cuda/dsa_attn.cu b/ggml/src/ggml-cuda/dsa_attn.cu
new file mode 100644
index 000000000..ba4639601
--- /dev/null
+++ b/ggml/src/ggml-cuda/dsa_attn.cu
@@ -0,0 +1,312 @@
+#include "dsa_attn.cuh"
+
+static inline bool v_is_k_view(const ggml_tensor * K, const ggml_tensor * V) {
+    if (!V || !V->data) return false;
+    auto k_data = (const char *)K->data;
+    auto v_data = (const char *)V->data;
+    auto k_row_size = ggml_row_size(K->type, K->ne[0]);
+    auto v_row_size = ggml_row_size(V->type, V->ne[0]);
+    return v_data >= k_data && v_data + v_row_size <= k_data + k_row_size;
+}
+
+static __global__ void k_prepare_mask(int nidx, const int * __restrict__ idx, const half * __restrict__ m_in,
+        half * __restrict__ m_out, size_t stride_idx, size_t stride_m) {
+    int row = blockIdx.x;
+    int col = blockIdx.y*blockDim.x + threadIdx.x;
+    idx += row*stride_idx;
+    m_out[row*nidx + col] = m_in[row*stride_m + idx[col]];
+}
+
+static __global__ void k_prepare_one_batch_kv(int nk, int ncol, const int * idx, const char * k_in,
+        half * k_out, size_t stride_k, size_t stride_idx) {
+    int row = blockIdx.y;
+    int col = blockIdx.x;
+    int i = idx[row*stride_idx + col];
+    auto k_row = (const half *)(k_in + stride_k * i);
+    k_out += (row*ncol + col)*nk;
+    for (int j = threadIdx.x; j < nk; j += blockDim.x) {
+        k_out[j] = k_row[j];
+    }
+}
+
+static __global__ void k_prepare_one_batch_q(int ne0, int ne1, size_t nb1, size_t nb2,
+        const float * q_in, half * q_out) {
+    int i0 = blockIdx.x*blockDim.x + threadIdx.x;
+    if (i0 >= ne0) {
+        return;
+    }
+    int i1 = blockIdx.y;
+    int i2 = blockIdx.z;
+    q_out[i0 + (i2 + i1*ne1)*ne0] = __float2half(q_in[i0 + i1*nb1 + i2*nb2]);
+}
+
+static __global__ void k_copy_dst(int nelem, const half * kqv16, float * dst) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= nelem) {
+        return;
+    }
+    dst[i] = __half2float(kqv16[i]);
+}
+
+template <int ncols_template, int block_size_template>
+static __global__ void soft_max_f16_simple(half * x, const half * mask, const int ncols_par, const int nrows_y, const float scale) {
+    const int ncols = ncols_template == 0 ? ncols_par : ncols_template;
+
+    const int tid  = threadIdx.x;
+    const int rowx = blockIdx.x;
+    const int rowy = rowx / nrows_y; // broadcast the mask in the row dimension
+
+    const int block_size = block_size_template == 0 ? blockDim.x : block_size_template;
+
+    const int warp_id = threadIdx.x / WARP_SIZE;
+    const int lane_id = threadIdx.x % WARP_SIZE;
+
+    extern __shared__ float data_soft_max_f32[];
+    float * buf_iw = data_soft_max_f32; // shared memory buffer for inter-warp communication
+    // shared memory buffer to cache values between iterations:
+    float * vals = buf_iw + WARP_SIZE;
+
+    float max_val = -INFINITY;
+
+#pragma unroll
+    for (int col0 = 0; col0 < ncols; col0 += block_size) {
+        const int col = col0 + tid;
+
+        if (ncols_template == 0 && col >= ncols) {
+            break;
+        }
+
+        const int64_t ix = (int64_t)rowx*ncols + col;
+        const int64_t iy = (int64_t)rowy*ncols + col;
+
+        const float val = scale*__half2float(x[ix]) + __half2float(mask[iy]);
+
+        vals[col] = val;
+        max_val = max(max_val, val);
+    }
+
+    // find the max value in the block
+    max_val = warp_reduce_max(max_val);
+    if (block_size > WARP_SIZE) {
+        if (warp_id == 0) {
+            buf_iw[lane_id] = -INFINITY;
+        }
+        __syncthreads();
+
+        if (lane_id == 0) {
+            buf_iw[warp_id] = max_val;
+        }
+        __syncthreads();
+
+        max_val = buf_iw[lane_id];
+        max_val = warp_reduce_max(max_val);
+    }
+
+    float tmp = 0.0f; // partial sum
+
+#pragma unroll
+    for (int col0 = 0; col0 < ncols; col0 += block_size) {
+        const int col = col0 + tid;
+
+        if (ncols_template == 0 && col >= ncols) {
+            break;
+        }
+
+        const float val = expf(vals[col] - max_val);
+        tmp += val;
+        vals[col] = val;
+    }
+
+    // find the sum of exps in the block
+    tmp = warp_reduce_sum(tmp);
+    if (block_size > WARP_SIZE) {
+        __syncthreads();
+        if (warp_id == 0) {
+            buf_iw[lane_id] = 0.0f;
+        }
+        __syncthreads();
+
+        if (lane_id == 0) {
+            buf_iw[warp_id] = tmp;
+        }
+        __syncthreads();
+
+        tmp = buf_iw[lane_id];
+        tmp = warp_reduce_sum(tmp);
+    }
+
+    const float inv_sum = 1.0f / tmp;
+
+#pragma unroll
+    for (int col0 = 0; col0 < ncols; col0 += block_size) {
+        const int col = col0 + tid;
+
+        if (ncols_template == 0 && col >= ncols) {
+            return;
+        }
+
+        const int64_t ix = (int64_t)rowx*ncols + col;
+        x[ix] = __float2half(vals[col] * inv_sum);
+    }
+}
+
+#define CUDA_SOFT_MAX_BLOCK_SIZE 1024
+
+static void soft_max_f16_cuda_simple(half * x, const half * mask, const int ncols_x, const int nrows_x,
+        const int nrows_y, const float scale, cudaStream_t stream) {
+    int nth = WARP_SIZE;
+    while (nth < ncols_x && nth < CUDA_SOFT_MAX_BLOCK_SIZE) nth *= 2;
+    const dim3 block_dims(nth,     1, 1);
+    const dim3 block_nums(nrows_x, 1, 1);
+    const size_t shmem = (GGML_PAD(ncols_x, WARP_SIZE) + WARP_SIZE)*sizeof(float);
+    static_assert(CUDA_SOFT_MAX_BLOCK_SIZE == 1024, "These values need to be adjusted.");
+
+    GGML_ASSERT(shmem < ggml_cuda_info().devices[ggml_cuda_get_device()].smpb);
+
+    switch (ncols_x) {
+        case 32:
+            soft_max_f16_simple<32, 32><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 64:
+            soft_max_f16_simple<64, 64><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 128:
+            soft_max_f16_simple<128, 128><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 256:
+            soft_max_f16_simple<256, 256><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 512:
+            soft_max_f16_simple<512, 512><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 1024:
+            soft_max_f16_simple<1024, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 2048:
+            soft_max_f16_simple<2048, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        case 4096:
+            soft_max_f16_simple<4096, 1024><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+        default:
+            soft_max_f16_simple<0, 0><<<block_nums, block_dims, shmem, stream>>>(x, mask, ncols_x, nrows_y, scale);
+            break;
+    }
+}
+
+bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    if (!dst) return false;
+
+    constexpr int k_max_rows = 32;
+
+    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 ggml_tensor * sink = dst->src[4];
+    const ggml_tensor * indexer = dst->src[5];
+
+    if (sink) return false; // We do not support sinks at this point
+    if (!Q || !K || !V || !mask || !indexer) return false;
+
+    if (indexer->ne[0] % 256 != 0) return false; // lazyness to add checks and handle tailes in case of not multiple of 256
+                                                 // But are there DSA variants where top_k is not a multiple of 256?
+    if (K->ne[1] < 4*indexer->ne[0]) return false; // for efficiency
+    if (K->ne[2] > 1 || K->ne[3] > 1 || mask->ne[2] > 1 || mask->ne[3] > 1 || Q->ne[3] > 1) return false;
+    if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16 || mask->type != GGML_TYPE_F16 || Q->type != GGML_TYPE_F32) return false;
+    if (K->ne[0] != Q->ne[0]) return false;
+
+    float scale;
+    memcpy(&scale, dst->op_params, sizeof(float));
+
+    const half alpha = 1.0f;
+    const half beta  = 0.0f;
+
+    int max_rows = std::min<int>(Q->ne[1], k_max_rows);
+    bool is_k_view = v_is_k_view(K, V);
+    auto mask_size = indexer->ne[0]*Q->ne[1]; // mask is relatively small, so we can do it once for the whole calculation
+    auto k_cache_size = indexer->ne[0]*K->ne[0]*max_rows;
+    auto v_cache_size = indexer->ne[0]*V->ne[0]*max_rows;
+    auto q_size = Q->ne[0]*Q->ne[2]*max_rows;
+    auto kq_size = indexer->ne[0]*Q->ne[2]*max_rows;
+    auto kqv_size = V->ne[0]*Q->ne[2]*max_rows;
+    ggml_cuda_pool_alloc<half> q16(ctx.pool(), q_size);
+    ggml_cuda_pool_alloc<half> kq16(ctx.pool(), kq_size);
+    ggml_cuda_pool_alloc<half> kqv16(ctx.pool(), kqv_size);
+    ggml_cuda_pool_alloc<half> mask16(ctx.pool(), mask_size);
+    ggml_cuda_pool_alloc<half> k16(ctx.pool(), k_cache_size);
+    ggml_cuda_pool_alloc<half> v16(ctx.pool());
+    size_t v_offset = 0;
+    if (is_k_view) {
+        v_offset = (const half *)V->data - (const half *)K->data;
+    } else {
+        v16.alloc(v_cache_size);
+    }
+    auto stride_idx = indexer->nb[1]/sizeof(int);
+    {
+        dim3 grid(Q->ne[1], indexer->ne[0]/256, 1);
+        k_prepare_mask<<<grid, 256, 0, ctx.stream()>>>(indexer->ne[0], (const int * )indexer->data,
+                (const half *)mask->data, mask16.get(), stride_idx, mask->nb[1]/sizeof(half));
+    }
+
+    int nstep = (Q->ne[1] + max_rows - 1)/max_rows;
+
+    for (int istep = 0; istep < nstep; ++istep) {
+        int first = istep*max_rows;
+        int last  = std::min<int>(first + max_rows, Q->ne[1]);
+        int nrows = last - first;
+        {
+            dim3 grid(indexer->ne[0], nrows, 1);
+            k_prepare_one_batch_kv<<<grid, 256, 0, ctx.stream()>>>(K->ne[0], indexer->ne[0],
+                    (const int *)indexer->data + stride_idx*first,
+                    (const char *)K->data, k16.get(), K->nb[1], stride_idx);
+            if (!is_k_view) {
+                k_prepare_one_batch_kv<<<grid, 256, 0, ctx.stream()>>>(V->ne[0], indexer->ne[0],
+                        (const int *)indexer->data + stride_idx*first,
+                        (const char *)V->data, v16.get(), V->nb[1], stride_idx);
+            }
+        }
+        {
+            int nblock = (Q->ne[0] + 255)/256;
+            dim3 grid(nblock, nrows, Q->ne[2]);
+            k_prepare_one_batch_q<<<grid, 256, 0, ctx.stream()>>>(Q->ne[0], Q->ne[2],
+                    Q->nb[1]/sizeof(float), Q->nb[2]/sizeof(float),
+                    (const float *)((const char *)Q->data + first*Q->nb[1]), q16.get());
+        }
+
+        CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
+                    indexer->ne[0], Q->ne[2], Q->ne[0],
+                    &alpha, k16.get(), K->ne[0], K->ne[0]*indexer->ne[0],
+                    q16.get(), Q->ne[0], Q->ne[0]*Q->ne[2],
+                    &beta, kq16.get(), indexer->ne[0], indexer->ne[0]*Q->ne[2], nrows));
+
+        soft_max_f16_cuda_simple(kq16.get(), mask16.get() + first*indexer->ne[0], indexer->ne[0], Q->ne[2]*nrows,
+                Q->ne[2], scale, ctx.stream());
+        CUDA_CHECK(cudaGetLastError());
+
+        if (is_k_view) {
+            CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N,
+                        V->ne[0], Q->ne[2], indexer->ne[0],
+                        &alpha, k16.get() + v_offset, K->ne[0], K->ne[0]*indexer->ne[0],
+                        kq16.get(), indexer->ne[0], indexer->ne[0]*Q->ne[2],
+                        &beta, kqv16.get(), V->ne[0], V->ne[0]*Q->ne[2], nrows));
+        } else {
+            CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_N, CUBLAS_OP_N,
+                        V->ne[0], Q->ne[2], indexer->ne[0],
+                        &alpha, v16.get(), V->ne[0], V->ne[0]*indexer->ne[0],
+                        kq16.get(), indexer->ne[0], indexer->ne[0]*Q->ne[2],
+                        &beta, kqv16.get(), V->ne[0], V->ne[0]*Q->ne[2], nrows));
+        }
+
+        {
+            int nelem = V->ne[0]*Q->ne[2]*nrows;
+            int nblock = (nelem + 255)/256;
+            k_copy_dst<<<nblock, 256, 0, ctx.stream()>>>(nelem, kqv16.get(),
+                    (float *)((char *)dst->data + dst->nb[2]*first));
+
+        }
+
+    }
+
+    return true;
+}
diff --git a/ggml/src/ggml-cuda/dsa_attn.cuh b/ggml/src/ggml-cuda/dsa_attn.cuh
new file mode 100644
index 000000000..7fa2e8df0
--- /dev/null
+++ b/ggml/src/ggml-cuda/dsa_attn.cuh
@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 02b35c4d0..cc9666be2 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -13,7 +13,7 @@
 #include "fattn-mma-f16-interface.cuh"
 #include "fattn-new-mma.cuh"
 #include "fattn.cuh"
-#include "convert.cuh"
+#include "dsa_attn.cuh"
 
 #include <cstdint>
 
@@ -43,6 +43,12 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
     const int32_t precision = KQV->op_params[3];
     const int32_t n_swa = KQV->op_params[4];
 
+    if (dst->src[5]) {
+        if (ggml_cuda_dsa_attn_ext(ctx, dst)) {
+            return;
+        }
+    }
+
     ggml_tensor local_dst, Kl, Vl, Ml;
     if (n_swa > 0) {
         int ntokens = std::max(FATTN_KQ_STRIDE, int(Q->ne[1]));