llama: create new fused op for QSA mask (#30640)

* llama: create new fused op for QSA scoring

* enroll glm5.3-next too

* tests : initialize the remaining tensors in test_kpool_mask

The harness sentinels were left uninitialized, so leftover memory could read as NaN and fail the NaN check.

* ggml: generalize the k-pool mask op into ggml_top_k_mask

The k-pool mask op only covered the pooled QSA selection of qwen4exp
and glm5-next, while the DSA models and the DeepSeek V4 CSA build the
same top-k mask by hand with fill, set_rows and add.

ggml_top_k_mask(mask, top_k, score, blocks, extra) keeps the mask
cells of the live picks and the extra cells, and sets the rest to
-inf. Every row tensor is 4D in the layout of top_k, and the mask is
broadcast to those rows, so a selection can also be made per head.
score, blocks and extra are optional: without score every pick is
live, without blocks a pick is a cell, and extra holds the cells a
row keeps without a pick.

llm_graph_context::build_fused_top_k_mask views the row tensors in
the layout of kq_mask and registers the fused node. The k-pool
models, the DSA attention of build_attn and the DeepSeek V4 CSA mask
use it when the probe enables the op, and keep their current path on
backends without it.

---------

Co-authored-by: Pascal <admin@serveurperso.com>

common/jinja/lexer.cpp

@@ -1,7 +1,9 @@
 #include "lexer.h"
 #include "runtime.h"
+#include "unicode.h"
 
 #include <cctype>
+#include <charconv>
 #include <functional>
 #include <map>
 #include <string>
@@ -75,6 +77,24 @@ lexer_result lexer::tokenize(const std::string & source) {
                 }
                 // add escaped char
                 char escaped_char = src[pos++];
+                if (escaped_char == 'x') {
+                    if (pos + 1 >= src.size()) {
+                        throw lexer_exception("incomplete hexadecimal escape", source, pos);
+                    }
+
+                    unsigned int value = 0;
+                    const char * first = &src[pos];
+                    const char * last  = first + 2;
+
+                    auto [ptr, ec] = std::from_chars(first, last, value, 16);
+                    if (ec != std::errc{} || ptr != last) {
+                        throw lexer_exception("invalid hexadecimal escape", source, pos);
+                    }
+
+                    str += common_unicode_cpt_to_utf8(value);
+                    pos += 2;
+                    continue;
+                }
                 if (escape_chars.find(escaped_char) == escape_chars.end()) {
                     throw lexer_exception(std::string("unknown escape character \\") + escaped_char, source, pos);
                 }

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -585,6 +585,7 @@ extern "C" {
         GGML_OP_DSV4_HC_COMB,
         GGML_OP_DSV4_HC_PRE,
         GGML_OP_DSV4_HC_POST,
+        GGML_OP_TOP_K_MASK,
 
         GGML_OP_UNARY,
 
@@ -2732,6 +2733,26 @@ extern "C" {
             struct ggml_tensor  * post,
             struct ggml_tensor  * comb);
 
+    // top-k sparse attention mask: keep the mask cells of the live picks and the extra cells, set the rest to -inf
+    //
+    // mask:   [n_kv,    m1,  m2,  m3]  f16 or f32, broadcast to the rows of top_k
+    // top_k:  [n_top,   ne1, ne2, ne3] i32, picks of each row
+    // score:  [n_pick,  ne1, ne2, ne3] f32, can be NULL: a pick is live when its score is above -inf, else every pick is live
+    // blocks: [n_blk,   n_pick]        i32, can be NULL: the member cells of each pick, else a pick is a cell
+    // extra:  [n_extra, ne1, ne2, ne3] i32, can be NULL: cells of each row kept without a pick
+    // res:    [n_kv,    ne1, ne2, ne3] same type as mask
+    //
+    // res[c, i] = mask[c, i % m] when c is a cell of a live pick of row i or an extra cell of row i, else -inf
+    // picks and cells out of range are skipped
+    //
+    GGML_API struct ggml_tensor * ggml_top_k_mask(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * top_k,
+            struct ggml_tensor  * score,
+            struct ggml_tensor  * blocks,
+            struct ggml_tensor  * extra);
+
     // custom operators
 
     typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);

ggml/src/ggml-backend-meta.cpp

@@ -774,6 +774,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         return src_ss[0];
     };
 
+    auto handle_pool_1d = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
+            return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+        }
+        return src_ss[0];
+    };
+
     auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
         GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
 
@@ -990,12 +997,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_CONV_3D:
             case GGML_OP_CONV_2D_DW:
             case GGML_OP_CONV_TRANSPOSE_2D:
-            case GGML_OP_POOL_1D:
             case GGML_OP_POOL_2D:
             case GGML_OP_POOL_2D_BACK:
             case GGML_OP_UPSCALE: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
+            case GGML_OP_POOL_1D: {
+                split_state = handle_pool_1d(src_ss);
+            } break;
             case GGML_OP_PAD: {
                 split_state = handle_pad(src_ss);
             } break;
@@ -1046,7 +1055,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             } break;
             case GGML_OP_DSV4_HC_COMB:
             case GGML_OP_DSV4_HC_PRE:
-            case GGML_OP_DSV4_HC_POST: {
+            case GGML_OP_DSV4_HC_POST:
+            case GGML_OP_TOP_K_MASK: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
             case GGML_OP_UNARY: {

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2105,6 +2105,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_dsv4_hc_post(params, tensor);
             } break;
+        case GGML_OP_TOP_K_MASK:
+            {
+                ggml_compute_forward_top_k_mask(params, tensor);
+            } break;
         case GGML_OP_MAP_CUSTOM1:
             {
                 ggml_compute_forward_map_custom1(params, tensor);
@@ -2288,6 +2292,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_DSV4_HC_COMB:
         case GGML_OP_DSV4_HC_PRE:
         case GGML_OP_DSV4_HC_POST:
+        case GGML_OP_TOP_K_MASK:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/ops.cpp

@@ -11736,6 +11736,92 @@ void ggml_compute_forward_dsv4_hc_post(
     }
 }
 
+// ggml_compute_forward_top_k_mask
+
+void ggml_compute_forward_top_k_mask(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+
+    GGML_ASSERT(dst->type == mask->type);
+    GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
+
+    const int64_t n_kv    = mask->ne[0];
+    const int64_t ne1     = dst->ne[1];
+    const int64_t ne2     = dst->ne[2];
+    const int64_t n_rows  = ggml_nrows(dst);
+    const int64_t n_top   = top_k->ne[0];
+    const int64_t n_blk   = blocks ? blocks->ne[0] : 1;
+    const int64_t n_pick  = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+    const int64_t n_extra = extra ? extra->ne[0] : 0;
+
+    const size_t ts = ggml_type_size(dst->type);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    for (int64_t ir = ith; ir < n_rows; ir += nth) {
+        const int64_t i1 = ir % ne1;
+        const int64_t i2 = (ir / ne1) % ne2;
+        const int64_t i3 = ir / (ne1*ne2);
+
+        const char * mask_row = (const char *) mask->data
+            + (i1 % mask->ne[1])*mask->nb[1] + (i2 % mask->ne[2])*mask->nb[2] + (i3 % mask->ne[3])*mask->nb[3];
+        char * dst_row = (char *) dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
+
+        if (dst->type == GGML_TYPE_F16) {
+            const ggml_fp16_t ninf = GGML_CPU_FP32_TO_FP16(-INFINITY);
+            for (int64_t c = 0; c < n_kv; ++c) {
+                ((ggml_fp16_t *) dst_row)[c] = ninf;
+            }
+        } else {
+            for (int64_t c = 0; c < n_kv; ++c) {
+                ((float *) dst_row)[c] = -INFINITY;
+            }
+        }
+
+        auto keep = [&](int32_t c) {
+            if (c >= 0 && c < n_kv) {
+                memcpy(dst_row + c*ts, mask_row + c*ts, ts);
+            }
+        };
+
+        const char * top_k_row = (const char *) top_k->data + i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3];
+
+        for (int64_t j = 0; j < n_top; ++j) {
+            const int32_t p = *(const int32_t *) (top_k_row + j*top_k->nb[0]);
+            if (p < 0 || p >= n_pick) {
+                continue;
+            }
+            if (score) {
+                const float s = *(const float *) ((const char *) score->data
+                    + p*score->nb[0] + i1*score->nb[1] + i2*score->nb[2] + i3*score->nb[3]);
+                if (!(s > -INFINITY)) {
+                    continue;
+                }
+            }
+            if (!blocks) {
+                keep(p);
+                continue;
+            }
+            for (int64_t k = 0; k < n_blk; ++k) {
+                keep(*(const int32_t *) ((const char *) blocks->data + k*blocks->nb[0] + p*blocks->nb[1]));
+            }
+        }
+
+        if (extra) {
+            const char * extra_row = (const char *) extra->data + i1*extra->nb[1] + i2*extra->nb[2] + i3*extra->nb[3];
+            for (int64_t j = 0; j < n_extra; ++j) {
+                keep(*(const int32_t *) (extra_row + j*extra->nb[0]));
+            }
+        }
+    }
+}
+
 // ggml_compute_forward_rwkv_wkv7
 
 static void ggml_compute_forward_rwkv_wkv7_f32(

ggml/src/ggml-cpu/ops.h

@@ -108,6 +108,7 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p
 void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_top_k_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cpu/vec.cpp

@@ -460,6 +460,92 @@ void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float *
     }
 }
 
+#ifndef GGML_GELU_QUICK_FP16
+// computed in f32 without the fp16 table, which loses precision on f32 input
+void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
+    int i = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        _mm512_storeu_ps(y + i, ggml_v_gelu_quick(_mm512_loadu_ps(x + i)));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        _mm256_storeu_ps(y + i, ggml_v_gelu_quick(_mm256_loadu_ps(x + i)));
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        _mm_storeu_ps(y + i, ggml_v_gelu_quick(_mm_loadu_ps(x + i)));
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svst1_f32(pg, y + i, ggml_v_gelu_quick(pg, svld1_f32(pg, x + i)));
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        vst1q_f32(y + i, ggml_v_gelu_quick(vld1q_f32(x + i)));
+    }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        vec_xst(ggml_v_gelu_quick(vec_xl(0, x + i)), 0, y + i);
+    }
+#elif defined(__riscv_v_intrinsic)
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
+        vfloat32m2_t vy = ggml_v_gelu_quick_m2(vx, vl);
+        __riscv_vse32_v_f32m2(&y[i], vy, vl);
+    }
+#endif
+    for (; i < n; ++i) {
+        y[i] = ggml_gelu_quick_f32(x[i]);
+    }
+}
+
+void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
+    int i = 0;
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+    for (; i + 15 < n; i += 16) {
+        _mm512_storeu_ps(y + i, _mm512_mul_ps(ggml_v_gelu_quick(_mm512_loadu_ps(x + i)), _mm512_loadu_ps(g + i)));
+    }
+#elif defined(__AVX2__) && defined(__FMA__)
+    for (; i + 7 < n; i += 8) {
+        _mm256_storeu_ps(y + i, _mm256_mul_ps(ggml_v_gelu_quick(_mm256_loadu_ps(x + i)), _mm256_loadu_ps(g + i)));
+    }
+#elif defined(__SSE2__)
+    for (; i + 3 < n; i += 4) {
+        _mm_storeu_ps(y + i, _mm_mul_ps(ggml_v_gelu_quick(_mm_loadu_ps(x + i)), _mm_loadu_ps(g + i)));
+    }
+#elif defined(__ARM_FEATURE_SVE) && defined(__aarch64__)
+    const int vlen = svcntw();
+    for (; i < n; i += vlen) {
+        const svbool_t pg = svwhilelt_b32_s32(i, n);
+        svst1_f32(pg, y + i, svmul_f32_x(pg, ggml_v_gelu_quick(pg, svld1_f32(pg, x + i)), svld1_f32(pg, g + i)));
+    }
+#elif defined(__ARM_NEON) && defined(__aarch64__)
+    for (; i + 3 < n; i += 4) {
+        vst1q_f32(y + i, vmulq_f32(ggml_v_gelu_quick(vld1q_f32(x + i)), vld1q_f32(g + i)));
+    }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        vec_xst(vec_mul(ggml_v_gelu_quick(vec_xl(0, x + i)), vec_xl(0, g + i)), 0, y + i);
+    }
+#elif defined(__riscv_v_intrinsic)
+    for (int vl; i < n; i += vl) {
+        vl = __riscv_vsetvl_e32m2(n - i);
+        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
+        vfloat32m2_t vg = __riscv_vle32_v_f32m2(&g[i], vl);
+        vfloat32m2_t vy = __riscv_vfmul_vv_f32m2(ggml_v_gelu_quick_m2(vx, vl), vg, vl);
+        __riscv_vse32_v_f32m2(&y[i], vy, vl);
+    }
+#endif
+    for (; i < n; ++i) {
+        y[i] = ggml_gelu_quick_f32(x[i]) * g[i];
+    }
+}
+#endif
+
 ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean) {
     int i = 0;
     ggml_float sum = 0;

ggml/src/ggml-cpu/vec.h

@@ -44,7 +44,6 @@ inline static ggml_float ggml_sve_sum_f32x2(svfloat32_t sum_lo, svfloat32_t sum_
 #endif
 
 #define GGML_GELU_FP16
-#define GGML_GELU_QUICK_FP16
 
 #define GGML_SOFT_MAX_UNROLL 4
 #define GGML_VEC_DOT_UNROLL  2
@@ -1054,11 +1053,7 @@ inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float *
     }
 }
 #else
-inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = ggml_gelu_quick_f32(x[i]);
-    }
-}
+void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x);
 #endif
 
 // Sigmoid Linear Unit (SiLU) function
@@ -1143,6 +1138,13 @@ inline static svfloat32_t ggml_v_silu(svbool_t pg, svfloat32_t x) {
     return svdiv_f32_x(pg, x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static svfloat32_t ggml_v_gelu_quick(svbool_t pg, svfloat32_t x) {
+    const svfloat32_t one = svdup_n_f32_x(pg, 1.0f);
+    const svfloat32_t exp_cx = ggml_v_expf(pg, svmul_n_f32_x(pg, x, GELU_QUICK_COEF));
+    return svdiv_f32_x(pg, x, svadd_f32_x(pg, one, exp_cx));
+}
+
 #elif defined(__ARM_NEON) && defined(__aarch64__)
 
 // adapted from arm limited optimized routine
@@ -1182,6 +1184,13 @@ inline static float32x4_t ggml_v_silu(float32x4_t x) {
     return vdivq_f32(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
+    const float32x4_t one = vdupq_n_f32(1.0f);
+    const float32x4_t exp_cx = ggml_v_expf(vmulq_n_f32(x, GELU_QUICK_COEF));
+    return vdivq_f32(x, vaddq_f32(one, exp_cx));
+}
+
 #elif defined(__AVX512F__) && defined(__AVX512DQ__)
 
 // adapted from arm limited optimized routine
@@ -1225,6 +1234,13 @@ inline static __m512 ggml_v_silu(__m512 x) {
     return _mm512_div_ps(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m512 ggml_v_gelu_quick(__m512 x) {
+    const __m512 one = _mm512_set1_ps(1);
+    const __m512 exp_cx = ggml_v_expf(_mm512_mul_ps(x, _mm512_set1_ps(GELU_QUICK_COEF)));
+    return _mm512_div_ps(x, _mm512_add_ps(one, exp_cx));
+}
+
 #elif defined(__AVX2__) && defined(__FMA__)
 
 // adapted from arm limited optimized routine
@@ -1280,6 +1296,13 @@ inline static __m256 ggml_v_silu(__m256 x) {
     return _mm256_div_ps(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m256 ggml_v_gelu_quick(__m256 x) {
+    const __m256 one = _mm256_set1_ps(1);
+    const __m256 exp_cx = ggml_v_expf(_mm256_mul_ps(x, _mm256_set1_ps(GELU_QUICK_COEF)));
+    return _mm256_div_ps(x, _mm256_add_ps(one, exp_cx));
+}
+
 #elif defined(__SSE2__) // __AVX2__ / __ARM_NEON
 
 #if defined(__FMA__)
@@ -1334,6 +1357,13 @@ inline static __m128 ggml_v_silu(__m128 x) {
     return _mm_div_ps(x, one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static __m128 ggml_v_gelu_quick(__m128 x) {
+    const __m128 one = _mm_set1_ps(1);
+    const __m128 exp_cx = ggml_v_expf(_mm_mul_ps(x, _mm_set1_ps(GELU_QUICK_COEF)));
+    return _mm_div_ps(x, _mm_add_ps(one, exp_cx));
+}
+
 #elif defined(__VXE__) || defined(__VXE2__)
 
 // adapted from arm limited optimized routine
@@ -1374,6 +1404,13 @@ inline static float32x4_t ggml_v_silu(float32x4_t v_x) {
     return vec_div(v_x, v_one_plus_exp_neg_x);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
+    const float32x4_t one = vec_splats(1.0f);
+    const float32x4_t exp_cx = ggml_v_expf(vec_mul(x, vec_splats(GELU_QUICK_COEF)));
+    return vec_div(x, vec_add(one, exp_cx));
+}
+
 #elif defined(__riscv_v_intrinsic)
 
 // adapted from arm limited optimized routine
@@ -1426,6 +1463,12 @@ inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) {
     return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl);
 }
 
+// gelu_quick(x) = x / (1 + exp(-1.702 x)); x stays the numerator so it cannot overflow
+inline static vfloat32m2_t ggml_v_gelu_quick_m2(vfloat32m2_t x, int vl) {
+    const vfloat32m2_t exp_cx = ggml_v_expf_m2(__riscv_vfmul_vf_f32m2(x, GELU_QUICK_COEF, vl), vl);
+    return __riscv_vfdiv_vv_f32m2(x, __riscv_vfadd_vf_f32m2(exp_cx, 1.0f, vl), vl);
+}
+
 #endif // __ARM_NEON / __AVX2__ / __SSE2__ / __VXE__ / __riscv_v_intrinsic
 
 inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
@@ -1571,11 +1614,7 @@ inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float
     }
 }
 #else
-inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = ggml_gelu_quick_f32(x[i]) * g[i];
-    }
-}
+void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g);
 #endif
 
 inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -70,6 +70,7 @@
 #include "ggml-cuda/cumsum.cuh"
 #include "ggml-cuda/fill.cuh"
 #include "ggml-cuda/lightning-indexer.cuh"
+#include "ggml-cuda/top-k-mask.cuh"
 #include "ggml.h"
 
 #include <algorithm>
@@ -2570,6 +2571,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_LIGHTNING_INDEXER:
             ggml_cuda_lightning_indexer(ctx, dst);
             break;
+        case GGML_OP_TOP_K_MASK:
+            ggml_cuda_op_top_k_mask(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -6003,6 +6007,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return true;
         case GGML_OP_LIGHTNING_INDEXER:
             return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
+        case GGML_OP_TOP_K_MASK:
+            return op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32;
 
         default:
             return false;

ggml/src/ggml-cuda/top-k-mask.cu

@@ -0,0 +1,141 @@
+#include "top-k-mask.cuh"
+#include "convert.cuh"
+
+#define CUDA_TOP_K_MASK_BLOCK_SIZE 512
+
+// one block per row: fill the row with -inf, then copy the mask cells of the live picks and the extra cells
+template <typename T>
+static __global__ void top_k_mask_kernel(
+        const T * mask, const int32_t * top_k, const float * score, const int32_t * blocks, const int32_t * extra, T * dst,
+        const int64_t n_kv, const int64_t ne1, const int64_t ne2,
+        const int64_t m1, const int64_t m2, const int64_t m3,
+        const int64_t n_top, const int64_t n_pick, const int64_t n_blk, const int64_t n_extra,
+        const int64_t sm1, const int64_t sm2, const int64_t sm3,
+        const int64_t sd1, const int64_t sd2, const int64_t sd3,
+        const int64_t sk0, const int64_t sk1, const int64_t sk2, const int64_t sk3,
+        const int64_t ss0, const int64_t ss1, const int64_t ss2, const int64_t ss3,
+        const int64_t sb0, const int64_t sb1,
+        const int64_t se0, const int64_t se1, const int64_t se2, const int64_t se3,
+        const bool vec) {
+    ggml_cuda_pdl_lc();
+
+    const int64_t ir = blockIdx.x;
+    const int64_t i1 = ir % ne1;
+    const int64_t i2 = (ir / ne1) % ne2;
+    const int64_t i3 = ir / (ne1*ne2);
+
+    const T       * mask_row  = mask  + (i1 % m1)*sm1 + (i2 % m2)*sm2 + (i3 % m3)*sm3;
+    T             * dst_row   = dst   + i1*sd1 + i2*sd2 + i3*sd3;
+    const int32_t * top_k_row = top_k + i1*sk1 + i2*sk2 + i3*sk3;
+
+    ggml_cuda_pdl_sync();
+
+    const T ninf = ggml_cuda_cast<T>(-INFINITY);
+    if (vec) {
+        constexpr int nv = sizeof(int4)/sizeof(T);
+        T fill[nv];
+#pragma unroll
+        for (int j = 0; j < nv; ++j) {
+            fill[j] = ninf;
+        }
+        const int4 fill4 = *(const int4 *) fill;
+        for (int64_t i = threadIdx.x; i < n_kv/nv; i += blockDim.x) {
+            ((int4 *) dst_row)[i] = fill4;
+        }
+    } else {
+        for (int64_t i = threadIdx.x; i < n_kv; i += blockDim.x) {
+            dst_row[i] = ninf;
+        }
+    }
+
+    __syncthreads();
+
+    // one thread per pick or extra cell; cells kept twice get the same value
+    for (int64_t i = threadIdx.x; i < n_top + n_extra; i += blockDim.x) {
+        if (i >= n_top) {
+            const int32_t c = extra[(i - n_top)*se0 + i1*se1 + i2*se2 + i3*se3];
+            if (c >= 0 && c < n_kv) {
+                dst_row[c] = mask_row[c];
+            }
+            continue;
+        }
+
+        const int32_t p = top_k_row[i*sk0];
+        if (p < 0 || p >= n_pick) {
+            continue;
+        }
+        if (score && !(score[p*ss0 + i1*ss1 + i2*ss2 + i3*ss3] > -INFINITY)) {
+            continue;
+        }
+        if (!blocks) {
+            if (p < n_kv) {
+                dst_row[p] = mask_row[p];
+            }
+            continue;
+        }
+        for (int64_t k = 0; k < n_blk; ++k) {
+            const int32_t c = blocks[k*sb0 + p*sb1];
+            if (c >= 0 && c < n_kv) {
+                dst_row[c] = mask_row[c];
+            }
+        }
+    }
+}
+
+template <typename T>
+static void top_k_mask_cuda(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * mask   = dst->src[0];
+    const ggml_tensor * top_k  = dst->src[1];
+    const ggml_tensor * score  = dst->src[2];
+    const ggml_tensor * blocks = dst->src[3];
+    const ggml_tensor * extra  = dst->src[4];
+
+    const int64_t n_kv = mask->ne[0];
+
+    GGML_ASSERT(dst->nb[0] == sizeof(T) && mask->nb[0] == sizeof(T));
+    GGML_ASSERT(dst->nb[1] % sizeof(T) == 0 && mask->nb[1] % sizeof(T) == 0);
+
+    // 16-byte stores need every dst row aligned
+    const bool vec = (n_kv*sizeof(T)) % sizeof(int4) == 0 &&
+        dst->nb[1] % sizeof(int4) == 0 && dst->nb[2] % sizeof(int4) == 0 && dst->nb[3] % sizeof(int4) == 0 &&
+        (uintptr_t) dst->data % sizeof(int4) == 0;
+
+    const int64_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+
+    const size_t si = sizeof(int32_t);
+    const size_t sf = sizeof(float);
+
+    const dim3 block_dims(CUDA_TOP_K_MASK_BLOCK_SIZE, 1, 1);
+    const dim3 grid_dims(ggml_nrows(dst), 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+    ggml_cuda_kernel_launch(top_k_mask_kernel<T>, launch_params,
+            (const T *) mask->data, (const int32_t *) top_k->data,
+            score  ? (const float   *) score->data  : nullptr,
+            blocks ? (const int32_t *) blocks->data : nullptr,
+            extra  ? (const int32_t *) extra->data  : nullptr,
+            (T *) dst->data,
+            n_kv, dst->ne[1], dst->ne[2],
+            mask->ne[1], mask->ne[2], mask->ne[3],
+            top_k->ne[0], n_pick, blocks ? blocks->ne[0] : 1, extra ? extra->ne[0] : 0,
+            mask->nb[1]/sizeof(T), mask->nb[2]/sizeof(T), mask->nb[3]/sizeof(T),
+            dst->nb[1]/sizeof(T),  dst->nb[2]/sizeof(T),  dst->nb[3]/sizeof(T),
+            top_k->nb[0]/si, top_k->nb[1]/si, top_k->nb[2]/si, top_k->nb[3]/si,
+            score  ? score->nb[0]/sf  : 0, score ? score->nb[1]/sf : 0, score ? score->nb[2]/sf : 0, score ? score->nb[3]/sf : 0,
+            blocks ? blocks->nb[0]/si : 0, blocks ? blocks->nb[1]/si : 0,
+            extra  ? extra->nb[0]/si  : 0, extra ? extra->nb[1]/si : 0, extra ? extra->nb[2]/si : 0, extra ? extra->nb[3]/si : 0,
+            vec);
+}
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    switch (dst->type) {
+        case GGML_TYPE_F16:
+            top_k_mask_cuda<half>(ctx, dst);
+            break;
+        case GGML_TYPE_F32:
+            top_k_mask_cuda<float>(ctx, dst);
+            break;
+        default:
+            GGML_ABORT("unsupported type");
+    }
+}

ggml/src/ggml-cuda/top-k-mask.cuh

@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_top_k_mask(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-hexagon/htp-repack.cpp

@@ -9,6 +9,11 @@
 #include <cmath>
 #include <cstring>
 #include <initializer_list>
+#include <thread>
+
+#if defined(__ARM_NEON) || defined(__aarch64__)
+#    include <arm_neon.h>
+#endif
 
 // --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) ---
 
@@ -43,8 +48,309 @@ static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const
     }
 }
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+static void repack_neon_transpose8x8(const uint8x8_t rows[8], uint8x8_t columns[8]) {
+    const uint8x8x2_t t0 = vtrn_u8(rows[0], rows[1]);
+    const uint8x8x2_t t1 = vtrn_u8(rows[2], rows[3]);
+    const uint8x8x2_t t2 = vtrn_u8(rows[4], rows[5]);
+    const uint8x8x2_t t3 = vtrn_u8(rows[6], rows[7]);
+    const uint16x4x2_t t4 = vtrn_u16(vreinterpret_u16_u8(t0.val[0]), vreinterpret_u16_u8(t1.val[0]));
+    const uint16x4x2_t t5 = vtrn_u16(vreinterpret_u16_u8(t0.val[1]), vreinterpret_u16_u8(t1.val[1]));
+    const uint16x4x2_t t6 = vtrn_u16(vreinterpret_u16_u8(t2.val[0]), vreinterpret_u16_u8(t3.val[0]));
+    const uint16x4x2_t t7 = vtrn_u16(vreinterpret_u16_u8(t2.val[1]), vreinterpret_u16_u8(t3.val[1]));
+    const uint32x2x2_t t8 = vtrn_u32(vreinterpret_u32_u16(t4.val[0]), vreinterpret_u32_u16(t6.val[0]));
+    const uint32x2x2_t t9 = vtrn_u32(vreinterpret_u32_u16(t4.val[1]), vreinterpret_u32_u16(t6.val[1]));
+    const uint32x2x2_t ta = vtrn_u32(vreinterpret_u32_u16(t5.val[0]), vreinterpret_u32_u16(t7.val[0]));
+    const uint32x2x2_t tb = vtrn_u32(vreinterpret_u32_u16(t5.val[1]), vreinterpret_u32_u16(t7.val[1]));
+    columns[0] = vreinterpret_u8_u32(t8.val[0]);
+    columns[1] = vreinterpret_u8_u32(ta.val[0]);
+    columns[2] = vreinterpret_u8_u32(t9.val[0]);
+    columns[3] = vreinterpret_u8_u32(tb.val[0]);
+    columns[4] = vreinterpret_u8_u32(t8.val[1]);
+    columns[5] = vreinterpret_u8_u32(ta.val[1]);
+    columns[6] = vreinterpret_u8_u32(t9.val[1]);
+    columns[7] = vreinterpret_u8_u32(tb.val[1]);
+}
+
+static uint8x8_t repack_neon_pack_4bit_pairs(uint8x8_t first, uint8x8_t second, bool high) {
+    const uint8x8x2_t separated = vuzp_u8(first, second);
+    const uint8x8_t mask = vdup_n_u8(0x0F);
+    uint8x8_t even = separated.val[0];
+    uint8x8_t odd  = separated.val[1];
+
+    if (high) {
+        even = vshr_n_u8(even, 4);
+        odd  = vshr_n_u8(odd, 4);
+    }
+
+    even = vand_u8(even, mask);
+    odd  = vand_u8(odd, mask);
+    return vorr_u8(vshl_n_u8(odd, 4), even);
+}
+
+static void repack_neon_4bit_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+    uint8x8_t low_rows[8];
+    uint8x8_t high_rows[8];
+    uint8x8_t columns[8];
+
+    for (int rr = 0; rr < 8; ++rr) {
+        const uint8x16_t packed = vld1q_u8(qs[rr]);
+        const uint8x8_t first = vget_low_u8(packed);
+        const uint8x8_t second = vget_high_u8(packed);
+        low_rows[rr]  = repack_neon_pack_4bit_pairs(first, second, false);
+        high_rows[rr] = repack_neon_pack_4bit_pairs(first, second, true);
+    }
+
+    repack_neon_transpose8x8(low_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+    }
+
+    repack_neon_transpose8x8(high_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+    }
+}
+
+static void repack_neon_k4_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0, bool high) {
+    uint8x8_t first_rows[8];
+    uint8x8_t second_rows[8];
+    uint8x8_t columns[8];
+
+    for (int rr = 0; rr < 8; ++rr) {
+        first_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr]), vld1_u8(qs[rr] + 8), high);
+        second_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr] + 16), vld1_u8(qs[rr] + 24), high);
+    }
+
+    repack_neon_transpose8x8(first_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+    }
+
+    repack_neon_transpose8x8(second_rows, columns);
+    for (int cp = 0; cp < 8; ++cp) {
+        vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+    }
+}
+
+static void repack_neon_q8_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+    uint8x8_t rows[8];
+    uint8x8_t columns[8];
+
+    for (int cp = 0; cp < 16; cp += 4) {
+        for (int rr = 0; rr < 8; ++rr) {
+            rows[rr] = vld1_u8(qs[rr] + 2 * cp);
+        }
+        repack_neon_transpose8x8(rows, columns);
+
+        for (int pair = 0; pair < 4; ++pair) {
+            const uint8x8x2_t interleaved = vzip_u8(columns[2 * pair], columns[2 * pair + 1]);
+            uint8_t * dst = tile_dst + (cp + pair) * 64 + 2 * row0;
+            vst1_u8(dst, interleaved.val[0]);
+            vst1_u8(dst + 8, interleaved.val[1]);
+        }
+    }
+}
+
+static inline void repack_store_u32(uint8_t * dst, uint32_t v) {
+    memcpy(dst, &v, sizeof(v));
+}
+
+static inline uint32x4_t repack_neon_transpose_2bit_4x4(uint32x4_t w) {
+    uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 6)), vdupq_n_u32(0x00CC00CCu));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 6)));
+    t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 12)), vdupq_n_u32(0x0000F0F0u));
+    return veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 12)));
+}
+
+static inline void repack_neon_q2_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+    const uint32x4_t w = repack_neon_transpose_2bit_4x4(vreinterpretq_u32_u8(vld1q_u8(qs)));
+    repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+    repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+    repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+    repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_q1_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+    uint32x4_t w = vreinterpretq_u32_u8(vld1q_u8(qs));
+    uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 1)), vdupq_n_u32(0x22222222u));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 1)));
+    t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 3)), vdupq_n_u32(0x0A0A0A0Au));
+    w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 3)));
+    w = repack_neon_transpose_2bit_4x4(w);
+    repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+    repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+    repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+    repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_pack_2bit_plane(uint8x16_t v0, uint8x16_t v1, uint8_t * dst, int row) {
+    const uint32x4x2_t z   = vuzpq_u32(vreinterpretq_u32_u8(v0), vreinterpretq_u32_u8(v1));
+    const uint32x4_t   p01 = vorrq_u32(z.val[0], vshlq_n_u32(z.val[1], 2));
+    const uint32x2x2_t u   = vuzp_u32(vget_low_u32(p01), vget_high_u32(p01));
+    const uint32x2_t   out = vorr_u32(u.val[0], vshl_n_u32(u.val[1], 4));
+    repack_store_u32(dst + 0 * 128 + row * 4, vget_lane_u32(out, 0));
+    repack_store_u32(dst + 1 * 128 + row * 4, vget_lane_u32(out, 1));
+}
+
+static inline void repack_neon_q2_3_k_low2(const uint8_t * qs, int kt_local, uint8_t * lo_pl, int row) {
+    const uint8_t *  qs_sub = qs + (kt_local / 4) * 32;
+    const int8x16_t  shift  = vdupq_n_s8(-2 * (kt_local % 4));
+    const uint8x16_t mask   = vdupq_n_u8(0x03);
+    const uint8x16_t v0     = vandq_u8(vshlq_u8(vld1q_u8(qs_sub +  0), shift), mask);
+    const uint8x16_t v1     = vandq_u8(vshlq_u8(vld1q_u8(qs_sub + 16), shift), mask);
+    repack_neon_pack_2bit_plane(v0, v1, lo_pl, row);
+}
+
+static inline void repack_neon_q3_k_hmask(const uint8_t * hmask, int kt_local, uint8_t * neg_pl, int row) {
+    static const uint32_t mask_words[4] = { 0x01010101u, 0x02020202u, 0x04040404u, 0x08080808u };
+    const uint8x16_t m0   = vreinterpretq_u8_u32(vld1q_u32(mask_words));
+    const uint8x16_t m1   = vshlq_n_u8(m0, 4);
+    const uint8x16_t hbit = vdupq_n_u8((uint8_t) (1u << kt_local));
+    const uint8x16_t b0   = vbicq_u8(m0, vtstq_u8(vld1q_u8(hmask +  0), hbit));
+    const uint8x16_t b1   = vbicq_u8(m1, vtstq_u8(vld1q_u8(hmask + 16), hbit));
+    const uint32x4_t b01  = vreinterpretq_u32_u8(vorrq_u8(b0, b1));
+    const uint32x2_t f0   = vorr_u32(vget_low_u32(b01), vget_high_u32(b01));
+    const uint32x2_t f1   = vpadd_u32(f0, f0);
+    repack_store_u32(neg_pl + row * 4, vget_lane_u32(f1, 0));
+}
+
+static inline void repack_neon_q6_k(const block_q6_K * b, int kt_local, uint8_t * lo_pl, uint8_t * hi_pl, int row) {
+    const int        c      = kt_local / 4;
+    const int        g_sb   = kt_local % 4;
+    const uint8_t *  ql     = b->ql + c * 64 + (g_sb & 1) * 32;
+    const uint8_t *  qh     = b->qh + c * 32;
+    const int8x16_t  lshift = vdupq_n_s8((g_sb >= 2) ? -4 : 0);
+    const int8x16_t  hshift = vdupq_n_s8(-2 * g_sb);
+    const uint8x16_t mask4  = vdupq_n_u8(0x0F);
+    const uint8x16_t mask2  = vdupq_n_u8(0x03);
+
+    const uint8x16_t lo0  = vandq_u8(vshlq_u8(vld1q_u8(ql +  0), lshift), mask4);
+    const uint8x16_t lo1  = vandq_u8(vshlq_u8(vld1q_u8(ql + 16), lshift), mask4);
+    const uint32x4x2_t lz = vuzpq_u32(vreinterpretq_u32_u8(lo0), vreinterpretq_u32_u8(lo1));
+    const uint32x4_t lout = vorrq_u32(lz.val[0], vshlq_n_u32(lz.val[1], 4));
+    repack_store_u32(lo_pl + 0 * 128 + row * 4, vgetq_lane_u32(lout, 0));
+    repack_store_u32(lo_pl + 1 * 128 + row * 4, vgetq_lane_u32(lout, 1));
+    repack_store_u32(lo_pl + 2 * 128 + row * 4, vgetq_lane_u32(lout, 2));
+    repack_store_u32(lo_pl + 3 * 128 + row * 4, vgetq_lane_u32(lout, 3));
+
+    const uint8x16_t hi0 = vandq_u8(vshlq_u8(vld1q_u8(qh +  0), hshift), mask2);
+    const uint8x16_t hi1 = vandq_u8(vshlq_u8(vld1q_u8(qh + 16), hshift), mask2);
+    repack_neon_pack_2bit_plane(hi0, hi1, hi_pl, row);
+}
+
+static inline void repack_neon_q5_k_qh(const uint8_t * qh, uint8_t hbit, uint8_t * plane, int row) {
+    static const uint8_t mask_bytes[16] = { 1, 2, 1, 2, 1, 2, 1, 2, 4, 8, 4, 8, 4, 8, 4, 8 };
+    const uint8x16_t m0   = vld1q_u8(mask_bytes);
+    const uint8x16_t m1   = vshlq_n_u8(m0, 4);
+    const uint8x16_t vbit = vdupq_n_u8(hbit);
+    const uint8x16_t b0   = vandq_u8(vtstq_u8(vld1q_u8(qh +  0), vbit), m0);
+    const uint8x16_t b1   = vandq_u8(vtstq_u8(vld1q_u8(qh + 16), vbit), m1);
+    const uint8x16_t b01  = vorrq_u8(b0, b1);
+    const uint8x8_t  fold = vorr_u8(vget_low_u8(b01), vget_high_u8(b01));
+    const uint8x8_t  out  = vpadd_u8(fold, fold);
+    plane[0 * 32 + row] = vget_lane_u8(out, 0);
+    plane[1 * 32 + row] = vget_lane_u8(out, 1);
+    plane[2 * 32 + row] = vget_lane_u8(out, 2);
+    plane[3 * 32 + row] = vget_lane_u8(out, 3);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const uint16_t * src) {
+    return vld1q_u16(src);
+}
+
+#if defined(__aarch64__) || defined(__ARM_FP16_FORMAT_IEEE)
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+    const float16x4_t lo = vcvt_f16_f32(vld1q_f32(src));
+    const float16x4_t hi = vcvt_f16_f32(vld1q_f32(src + 4));
+    return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+    const uint16x8_t  bf   = vld1q_u16((const uint16_t *) src);
+    const float32x4_t f_lo = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(bf), 16));
+    const float32x4_t f_hi = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(bf), 16));
+    const float16x4_t lo   = vcvt_f16_f32(f_lo);
+    const float16x4_t hi   = vcvt_f16_f32(f_hi);
+    return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+#else
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+    uint16_t tmp[8];
+    for (int i = 0; i < 8; ++i) {
+        tmp[i] = GGML_FP32_TO_FP16(src[i]);
+    }
+    return vld1q_u16(tmp);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+    uint16_t tmp[8];
+    for (int i = 0; i < 8; ++i) {
+        tmp[i] = GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(src[i]));
+    }
+    return vld1q_u16(tmp);
+}
+#endif
+
+template <typename SrcT>
+static void repack_neon_f16_tile(uint16_t * tile_dst, const SrcT * src_slice, int64_t row_stride, int ct, int kt) {
+    for (int row0 = 0; row0 < 32; row0 += 4) {
+        const SrcT * r0 = src_slice + (int64_t) (ct * 32 + row0 + 0) * row_stride + kt * 32;
+        const SrcT * r1 = src_slice + (int64_t) (ct * 32 + row0 + 1) * row_stride + kt * 32;
+        const SrcT * r2 = src_slice + (int64_t) (ct * 32 + row0 + 2) * row_stride + kt * 32;
+        const SrcT * r3 = src_slice + (int64_t) (ct * 32 + row0 + 3) * row_stride + kt * 32;
+        for (int cp = 0; cp < 16; cp += 4) {
+            const uint32x4_t v0 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r0 + 2 * cp));
+            const uint32x4_t v1 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r1 + 2 * cp));
+            const uint32x4_t v2 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r2 + 2 * cp));
+            const uint32x4_t v3 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r3 + 2 * cp));
+
+            const uint32x2x2_t t0 = vtrn_u32(vget_low_u32(v0),  vget_low_u32(v1));
+            const uint32x2x2_t t1 = vtrn_u32(vget_high_u32(v0), vget_high_u32(v1));
+            const uint32x2x2_t t2 = vtrn_u32(vget_low_u32(v2),  vget_low_u32(v3));
+            const uint32x2x2_t t3 = vtrn_u32(vget_high_u32(v2), vget_high_u32(v3));
+
+            vst1q_u16(tile_dst + (cp + 0) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[0], t2.val[0])));
+            vst1q_u16(tile_dst + (cp + 1) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[1], t2.val[1])));
+            vst1q_u16(tile_dst + (cp + 2) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[0], t3.val[0])));
+            vst1q_u16(tile_dst + (cp + 3) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[1], t3.val[1])));
+        }
+    }
+}
+#endif
+
 // --- Generic 32x32 tile grid walkers for block-32 and float types ---
 
+template <typename Fn>
+static void repack_parallel_for_ct(size_t slice_bytes, int n_col_tiles, Fn && fn) {
+    int n_threads = 1;
+    if (slice_bytes > 128 * 1024) {
+        for (int nt : { 4, 3, 2 }) {
+            if (n_col_tiles % nt == 0) {
+                n_threads = nt;
+                break;
+            }
+        }
+    }
+
+    if (n_threads <= 1) {
+        fn(0, n_col_tiles);
+        return;
+    }
+
+    const int chunk = n_col_tiles / n_threads;
+    std::thread workers[3];
+    for (int i = 0; i < n_threads - 1; ++i) {
+        workers[i] = std::thread([&fn, i, chunk]() {
+            fn(i * chunk, (i + 1) * chunk);
+        });
+    }
+    fn((n_threads - 1) * chunk, n_col_tiles);
+    for (int i = 0; i < n_threads - 1; ++i) {
+        workers[i].join();
+    }
+}
+
 template <typename SrcElemT, int ELEMS_PER_KTILE, typename PackTileFn>
 static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t offset, size_t size,
                                     size_t tile_size, PackTileFn && pack_tile) {
@@ -70,12 +376,14 @@ static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t o
         const SrcElemT * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * row_stride);
         uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
 
-        for (int ct = 0; ct < n_col_tiles; ct++) {
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            for (int ct = ct0; ct < ct1; ct++) {
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+                }
             }
-        }
+        });
     }
 }
 
@@ -130,8 +438,24 @@ static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t
 static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0,
         [](uint8_t * tile_dst, const block_q4_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 if (r < ne1 && kt < ne0 / 32) {
@@ -167,8 +491,26 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset,
 static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q4_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_1,
         [](uint8_t * tile_dst, const block_q4_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    const block_q4_1 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+                    scale_dst[2 * row + 0] = src_block.d;
+                    scale_dst[2 * row + 1] = src_block.m;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 if (r < ne1 && kt < ne0 / 32) {
@@ -207,8 +549,24 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset,
 static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4,
         [](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
-            uint8_t tile_quants[32][32];
             uint8_t * scale_dst = tile_dst + 512;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_4bit_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].e;
+                }
+                return;
+            }
+#endif
+            uint8_t tile_quants[32][32];
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 if (r < ne1 && kt < ne0 / 32) {
@@ -244,6 +602,23 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset
 static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_to_tiled_grid_32<block_q8_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q8_0,
         [](uint8_t * tile_dst, const block_q8_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+            ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+                for (int row0 = 0; row0 < 32; row0 += 8) {
+                    const uint8_t * qs[8];
+                    for (int rr = 0; rr < 8; ++rr) {
+                        qs[rr] = (const uint8_t *) src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+                    }
+                    repack_neon_q8_rows(qs, tile_dst, row0);
+                }
+
+                for (int row = 0; row < 32; ++row) {
+                    scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+                }
+                return;
+            }
+#endif
             for (int cp = 0; cp < 16; cp++) {
                 const int col0 = cp * 2;
                 const int col1 = col0 + 1;
@@ -254,7 +629,6 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset,
                     tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0;
                 }
             }
-            ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
             for (int row = 0; row < 32; row++) {
                 const int64_t r = ct * 32 + row;
                 scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * row_stride + kt].d : 0;
@@ -303,24 +677,33 @@ static void repack_q2_0_tiled(ggml_tensor * t, const void * data, size_t offset,
     for (int64_t slice = 0; slice < t->ne[2] * t->ne[3]; ++slice) {
         const block_q2_0 * src_slice = src + slice * t->ne[1] * blocks_per_row;
         uint8_t * dst_slice = (uint8_t *) t->data + slice * matrix_size;
-        memset(dst_slice, 0, matrix_size);
+        const size_t ct_bytes = (size_t) n_weight_tiles * HTP_TENSOR_TILE_SIZE_Q2_0;
 
-        for (int64_t row_index = 0; row_index < t->ne[1]; ++row_index) {
-            const int64_t ct = row_index / 32;
-            const int row = (int) (row_index % 32);
-            for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
-                const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
-                uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
-                for (int group = 0; group < 16; ++group) {
-                    for (int byte = 0; byte < 4; ++byte) {
-                        const int k = 4 * group + byte;
-                        const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
-                        tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+        repack_parallel_for_ct(matrix_size, (int) n_col_tiles, [&](int ct0, int ct1) {
+            memset(dst_slice + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
+            const int64_t r0 = (int64_t) ct0 * 32;
+            const int64_t r1 = (std::min)((int64_t) ct1 * 32, t->ne[1]);
+            for (int64_t row_index = r0; row_index < r1; ++row_index) {
+                const int64_t ct = row_index / 32;
+                const int row = (int) (row_index % 32);
+                for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
+                    const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
+                    uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+                    repack_neon_q2_0_block(block->qs, tile, row);
+#else
+                    for (int group = 0; group < 16; ++group) {
+                        for (int byte = 0; byte < 4; ++byte) {
+                            const int k = 4 * group + byte;
+                            const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
+                            tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+                        }
                     }
+#endif
+                    memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
                 }
-                memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
             }
-        }
+        });
     }
 }
 
@@ -368,6 +751,12 @@ static void repack_float_to_f16_tiled(ggml_tensor * t, const void * data, size_t
     repack_to_tiled_grid_32<SrcT, 32>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_F16,
         [&](uint8_t * tile_bytes, const SrcT * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
             uint16_t * tile_dst = (uint16_t *) tile_bytes;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            if (ct * 32 + 32 <= ne1 && kt * 32 + 32 <= ne0) {
+                repack_neon_f16_tile(tile_dst, src_slice, row_stride, ct, kt);
+                return;
+            }
+#endif
             for (int cp = 0; cp < 16; cp++) {
                 const int64_t c0 = kt * 32 + cp * 2;
                 const int64_t c1 = c0 + 1;
@@ -424,20 +813,25 @@ static void repack_kquant_to_tiled_grid(ggml_tensor * t, const void * data, size
     for (int64_t slice_idx = 0; slice_idx < ne2 * ne3; slice_idx++) {
         const BlockT * src_slice = src_matrix + slice_idx * (ne1 * sb_per_row);
         uint8_t * matrix_dst     = (uint8_t *) t->data + slice_idx * matrix_size;
+        const size_t ct_bytes    = (size_t) n_k_tiles * tile_size;
 
-        // Padding rows and OR-ed bit planes require zeroed tiles
-        memset(matrix_dst, 0, matrix_size);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            // Padding rows and OR-ed bit planes require zeroed tiles
+            memset(matrix_dst + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
 
-        for (int64_t r = 0; r < ne1; r++) {
-            const int ct  = (int) (r / 32);
-            const int row = (int) (r % 32);
-            const BlockT * src_row = src_slice + r * sb_per_row;
+            const int64_t r0 = (int64_t) ct0 * 32;
+            const int64_t r1 = (std::min)((int64_t) ct1 * 32, ne1);
+            for (int64_t r = r0; r < r1; r++) {
+                const int ct  = (int) (r / 32);
+                const int row = (int) (r % 32);
+                const BlockT * src_row = src_slice + r * sb_per_row;
 
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+                }
             }
-        }
+        });
     }
 }
 
@@ -493,6 +887,7 @@ static void unpack_kquant_from_tiled_grid(void * data, const ggml_tensor * t, si
 
 // --- K-Quant Helpers & Per-Type Implementations (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) ---
 
+#if !defined(__ARM_NEON) && !defined(__aarch64__)
 // unsigned 6-bit value (0..63) of element e of a Q6_K block, same bit layout as dequantize_row_q6_K
 static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
     const int c = e / 128;
@@ -510,6 +905,7 @@ static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
     }
     return (uint8_t) (lo | (hi << 4));
 }
+#endif
 
 static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
     repack_kquant_to_tiled_grid<block_q6_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q6_K,
@@ -519,6 +915,9 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t * hi_pl   = tile + 512;
             ggml_half * sc_pl = (ggml_half *) (tile + 768);
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q6_k(b, kt_local, lo_pl, hi_pl, row);
+#else
             for (int lk = 0; lk < 32; lk++) {
                 const uint8_t q6 = q6_K_get_quant(b, kt_local * 32 + lk);
                 const int g   = lk >> 2;
@@ -526,6 +925,7 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                 lo_pl[(g >> 1) * 128 + pos] |= (uint8_t) ((q6 & 0xF) << ((g & 1) * 4));
                 hi_pl[(g >> 2) * 128 + pos] |= (uint8_t) ((q6 >> 4) << ((g & 3) * 2));
             }
+#endif
             for (int sub = 0; sub < 2; sub++) {
                 sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) b->scales[kt_local * 2 + sub]);
             }
@@ -585,6 +985,7 @@ static void repack_tiled_q6_K(void * data, const ggml_tensor * t, size_t offset,
         });
 }
 
+#if !defined(__ARM_NEON) && !defined(__aarch64__)
 // low 2 bits (0..3) of element e of a Q2_K or Q3_K block, same bit layout as dequantize_row_q2_K / q3_K
 static inline uint8_t q2_3_K_get_low2(const uint8_t * qs, int e) {
     const int c = e / 128;
@@ -600,6 +1001,7 @@ static inline bool q3_K_get_hbit(const block_q3_K * b, int e) {
     const int l = e % 32;
     return (b->hmask[l] >> (c * 4 + j)) & 1;
 }
+#endif
 
 // signed 6-bit scale j (-32..31) of a Q3_K block, same packing as quantize_row_q3_K_ref
 static inline int q3_K_get_scale(const uint8_t * scales, int j) {
@@ -670,6 +1072,10 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t * neg_pl  = tile + 256;
             ggml_half * sc_pl = (ggml_half *) (tile + 384);
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+            repack_neon_q3_k_hmask(b->hmask, kt_local, neg_pl, row);
+#else
             for (int lk = 0; lk < 32; lk++) {
                 const int e   = kt_local * 32 + lk;
                 const int g   = lk >> 2;
@@ -679,6 +1085,7 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                     neg_pl[pos] |= (uint8_t) (1 << g);
                 }
             }
+#endif
             for (int sub = 0; sub < 2; sub++) {
                 sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) q3_K_get_scale(b->scales, kt_local * 2 + sub));
             }
@@ -750,11 +1157,15 @@ static void repack_q2_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t * sc_pl   = tile + 256;
             ggml_half * dd_pl = (ggml_half *) (tile + 320);
 
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+#else
             for (int lk = 0; lk < 32; lk++) {
                 const int g   = lk >> 2;
                 const int pos = row * 4 + (lk & 3);
                 lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, kt_local * 32 + lk) << ((g & 3) * 2));
             }
+#endif
             for (int sub = 0; sub < 2; sub++) {
                 sc_pl[sub * 32 + row] = b->scales[kt_local * 2 + sub];
             }
@@ -846,23 +1257,46 @@ static inline void pack_k4_scales_mins(uint8_t * scales, ggml_half * d, ggml_hal
     }
 }
 
+template <typename BlockT>
+static inline void repack_k4_quants(uint8_t * tile_dst, const BlockT * b, int row, int kt_local,
+                                    int64_t sb_per_row, bool use_neon) {
+#if defined(__ARM_NEON) || defined(__aarch64__)
+    if (use_neon) {
+        if ((row & 7) == 0) {
+            const uint8_t * qs[8];
+            for (int rr = 0; rr < 8; ++rr) {
+                qs[rr] = b[rr * sb_per_row].qs + (kt_local / 2) * 32;
+            }
+            repack_neon_k4_rows(qs, tile_dst, row, (kt_local & 1) != 0);
+        }
+        return;
+    }
+#else
+    GGML_UNUSED(sb_per_row);
+    GGML_UNUSED(use_neon);
+#endif
+    const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
+    const int shift = (kt_local & 1) ? 4 : 0;
+
+    for (int cp = 0; cp < 16; cp++) {
+        const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
+        const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
+        tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
+    }
+}
+
 static void repack_q4_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    const int64_t sb_per_row = t->ne[0] / QK_K;
+    const bool    use_neon   = (t->ne[1] % 8 == 0);
     repack_kquant_to_tiled_grid<block_q4_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_K,
-        [](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
+        [sb_per_row, use_neon](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
             const float d    = GGML_FP16_TO_FP32(b->d);
             const float dmin = GGML_FP16_TO_FP32(b->dmin);
 
             uint8_t sc, m;
             get_scale_min_k4(kt_local, b->scales, &sc, &m);
 
-            const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
-            const int shift = (kt_local & 1) ? 4 : 0;
-
-            for (int cp = 0; cp < 16; cp++) {
-                const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
-                const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
-                tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-            }
+            repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
 
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
             scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -900,8 +1334,10 @@ static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset,
 }
 
 static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+    const int64_t sb_per_row = t->ne[0] / QK_K;
+    const bool    use_neon   = (t->ne[1] % 8 == 0);
     repack_kquant_to_tiled_grid<block_q5_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_K,
-        [](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
+        [sb_per_row, use_neon](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
             const float d     = GGML_FP16_TO_FP32(b->d);
             const float dmin  = GGML_FP16_TO_FP32(b->dmin);
             uint8_t * plane   = tile_dst + 640;
@@ -909,15 +1345,13 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
             uint8_t sc, m;
             get_scale_min_k4(kt_local, b->scales, &sc, &m);
 
-            const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
-            const int shift    = (kt_local & 1) ? 4 : 0;
+            repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
+
             const uint8_t hbit = (uint8_t) (1 << kt_local);
-
+#if defined(__ARM_NEON) || defined(__aarch64__)
+            repack_neon_q5_k_qh(b->qh, hbit, plane, row);
+#else
             for (int cp = 0; cp < 16; cp++) {
-                const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
-                const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
-                tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-
                 const int i    = cp / 4;
                 const int lane = (cp % 4) * 32 + row;
                 if (b->qh[2 * cp + 0] & hbit) {
@@ -927,6 +1361,7 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
                     plane[lane] |= (uint8_t) (1 << (2 * i + 1));
                 }
             }
+#endif
 
             ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
             scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -1002,31 +1437,37 @@ static void repack_q1_0_tiled(ggml_tensor * t, const void * data, size_t offset,
         const block_q1_0 * src_slice = src_matrix + (slice_idx - start_slice) * ne1 * n_k_tiles;
         uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
 
-        for (int ct = 0; ct < n_col_tiles; ct++) {
-            for (int kt = 0; kt < n_k_tiles; kt++) {
-                uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
-                ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
-                memset(tile_dst, 0, tile_size);
+        repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+            for (int ct = ct0; ct < ct1; ct++) {
+                for (int kt = 0; kt < n_k_tiles; kt++) {
+                    uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+                    ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+                    memset(tile_dst, 0, tile_size);
 
-                for (int row = 0; row < 32; row++) {
-                    const int64_t r = ct * 32 + row;
-                    if (r >= ne1) {
-                        continue;
-                    }
-
-                    const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
-                    scale_dst[row] = block.d;
-                    for (int k = 0; k < QK1_0; k++) {
-                        int byte_offset;
-                        int bit;
-                        q1_0_tiled_bit_location(k, row, byte_offset, bit);
-                        if (block.qs[k / 8] & (1u << (k % 8))) {
-                            tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+                    for (int row = 0; row < 32; row++) {
+                        const int64_t r = ct * 32 + row;
+                        if (r >= ne1) {
+                            continue;
                         }
+
+                        const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
+                        scale_dst[row] = block.d;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+                        repack_neon_q1_0_block(block.qs, tile_dst, row);
+#else
+                        for (int k = 0; k < QK1_0; k++) {
+                            int byte_offset;
+                            int bit;
+                            q1_0_tiled_bit_location(k, row, byte_offset, bit);
+                            if (block.qs[k / 8] & (1u << (k % 8))) {
+                                tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+                            }
+                        }
+#endif
                     }
                 }
             }
-        }
+        });
     }
 }
 
@@ -1127,11 +1568,11 @@ void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t siz
             break;
         case GGML_TYPE_F32:
             repack_float_to_f16_tiled<float>(tensor, data, 0, size,
-                [](float v) { return (uint16_t) ggml_fp32_to_fp16(v); });
+                [](float v) { return (uint16_t) GGML_FP32_TO_FP16(v); });
             break;
         case GGML_TYPE_BF16:
             repack_float_to_f16_tiled<ggml_bf16_t>(tensor, data, 0, size,
-                [](ggml_bf16_t v) { return (uint16_t) ggml_fp32_to_fp16(ggml_bf16_to_fp32(v)); });
+                [](ggml_bf16_t v) { return (uint16_t) GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(v)); });
             break;
         default:
             break;

ggml/src/ggml-metal/ggml-metal-context.m

@@ -539,8 +539,14 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
                 descriptor.outputURL = [NSURL fileURLWithPath:path];
 
                 NSError * error = nil;
-                if (![[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error]) {
+                const bool ok = [[MTLCaptureManager sharedCaptureManager] startCaptureWithDescriptor:descriptor error:&error];
+
+                [descriptor release];
+
+                if (!ok) {
                     GGML_LOG_ERROR("%s: error: unable to start capture '%s' (did you set METAL_CAPTURE_ENABLED=1 ?)\n", __func__, [[error localizedDescription] UTF8String]);
+                    [ctx->capture_scope release];
+                    ctx->capture_scope = nil;
                 } else {
                     [ctx->capture_scope beginScope];
                     ctx->capture_started = true;
@@ -600,6 +606,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
         // enter here only when capturing in order to wait for all computation to finish
         // otherwise, we leave the graph to compute asynchronously
         if (use_capture && ctx->capture_started) {
+            enum ggml_status res = GGML_STATUS_SUCCESS;
+
             // wait for completion and check status of each command buffer
             // needed to detect if the device ran out-of-memory for example (#1881)
             {
@@ -613,11 +621,11 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
                         GGML_LOG_INFO("error: %s\n", [[cmd_buf error].localizedDescription UTF8String]);
                     }
 
-                    return GGML_STATUS_FAILED;
+                    res = GGML_STATUS_FAILED;
                 }
             }
 
-            for (int i = 0; i < n_cb; ++i) {
+            for (int i = 0; i < n_cb && res == GGML_STATUS_SUCCESS; ++i) {
                 id<MTLCommandBuffer> cmd_buf = ctx->cmd_bufs[i].obj;
                 [cmd_buf waitUntilCompleted];
 
@@ -627,8 +635,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
                     if (status == MTLCommandBufferStatusError) {
                         GGML_LOG_INFO("error: %s\n", [[cmd_buf error].localizedDescription UTF8String]);
                     }
-
-                    return GGML_STATUS_FAILED;
+                    res = GGML_STATUS_FAILED;
+                    break;
                 }
 
                 id<MTLCommandBuffer> next_buffer = (i + 1 < n_cb ? ctx->cmd_bufs[i + 1].obj : nil);
@@ -643,7 +651,8 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
 
                 if (ctx->abort_callback && ctx->abort_callback(ctx->abort_callback_data)) {
                     GGML_LOG_INFO("%s: command buffer %d aborted", __func__, i);
-                    return GGML_STATUS_ABORTED;
+                    res = GGML_STATUS_ABORTED;
+                    break;
                 }
 
                 [next_buffer commit];
@@ -652,7 +661,14 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
             [ctx->capture_scope endScope];
             [[MTLCaptureManager sharedCaptureManager] stopCapture];
 
+            [ctx->capture_scope release];
+            ctx->capture_scope = nil;
+
             ctx->capture_started = false;
+
+            if (res != GGML_STATUS_SUCCESS) {
+                return res;
+            }
         }
     }
 

ggml/src/ggml-opencl/CMakeLists.txt

@@ -80,6 +80,8 @@ set(GGML_OPENCL_KERNELS
     solve_tri
     im2col_f32
     im2col_f16
+    im2col_3d_f32
+    im2col_3d_f16
     mean
     mul_mv_f16_f16
     mul_mv_f16_f32_1row

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -971,6 +971,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_mul_mv_iq4_nl_f32_flat;
     cl_kernel kernel_solve_tri_f32;
     cl_kernel kernel_im2col_f32, kernel_im2col_f16;
+    cl_kernel kernel_im2col_3d_f32, kernel_im2col_3d_f16;
     cl_kernel kernel_argsort_f32_i32;
     cl_kernel kernel_sum_rows_f32, kernel_sum_rows_f32_4;
     cl_kernel kernel_cumsum_blk, kernel_cumsum_add;
@@ -1243,6 +1244,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_q4_1_f32;
     cl_kernel kernel_gemm_noshuffle_q8_0_f32, kernel_gemm_noshuffle_q8_0_f32_bin;
     cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a = nullptr;  // dp4a (int8) dense q8_0 prefill GEMM (opt-in)
+    cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg = nullptr;  // q8_0 dense dp4a, weights via texture (opt-in)
     cl_kernel kernel_gemm_noshuffle_q8_0_q8_1_dp4a_ila_a8_bin = nullptr;
     cl_kernel kernel_gemv_noshuffle_q8_0_f32;
@@ -1264,6 +1266,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_ila_a8_bin;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_32b_trans;
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = nullptr;  // dp4a dense prefill GEMM, weights via texture (X1 opt-in)
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow = nullptr;
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg = nullptr;
@@ -1271,7 +1274,9 @@ struct ggml_backend_opencl_context {
     int tile_n_dp4a_narrow_q4_k = 32;
     int max_n_dp4a_narrow_q4_k = 16;
     cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q5_K prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q6_K prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_quant_a_q8_1;                    // plain activation q8_1 pre-pass
     cl_kernel kernel_gemm_noshuffle_q4_k_f32_r1;
     cl_kernel kernel_gemm_noshuffle_q4_k_f32_kimg;
@@ -1305,6 +1310,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_iq4_nl_f32;
     cl_kernel kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a = nullptr;  // dp4a (int8) dense IQ4_NL prefill GEMM
     cl_kernel kernel_gemm_noshuffle_q4_0_q8_1_dp4a = nullptr;  // dp4a (int8) dense q4_0 prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4 = nullptr;
     cl_kernel kernel_gemm_coopk_q4_0_q8_1_dp4a_c4 = nullptr;
     cl_kernel kernel_gemm_coopk_q4_0_q8_1_dp4a_c2 = nullptr;
     int       nsg_coopk_dp4a_q4_0  = 8;
@@ -2096,6 +2102,40 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
+    // im2col_3d_f32
+    {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+        const std::string kernel_src {
+            #include "im2col_3d_f32.cl.h"
+        };
+#else
+        const std::string kernel_src = read_file("im2col_3d_f32.cl");
+#endif
+        cl_program prog =
+            build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+
+        CL_CHECK((backend_ctx->kernel_im2col_3d_f32 = clCreateKernel(prog, "kernel_im2col_3d_f32", &err), err));
+        CL_CHECK(clReleaseProgram(prog));
+        GGML_LOG_CONT(".");
+    }
+
+    // im2col_3d_f16
+    {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+        const std::string kernel_src {
+            #include "im2col_3d_f16.cl.h"
+        };
+#else
+        const std::string kernel_src = read_file("im2col_3d_f16.cl");
+#endif
+        cl_program prog =
+            build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+
+        CL_CHECK((backend_ctx->kernel_im2col_3d_f16 = clCreateKernel(prog, "kernel_im2col_3d_f16", &err), err));
+        CL_CHECK(clReleaseProgram(prog));
+        GGML_LOG_CONT(".");
+    }
+
     // mul_mv_q4_0_f32
     {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -4347,6 +4387,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #endif
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_0_q8_1_dp4a", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -4486,6 +4527,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), dp4a_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a", &err), err));
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
 
         // Compile a smaller N tile to avoid unused columns in small batches.
@@ -4518,6 +4560,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a", &err), err));
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -4533,6 +4576,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #endif
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_k_q8_1_dp4a", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -4548,6 +4592,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #endif
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q6_k_q8_1_dp4a", &err), err));
+        CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4 = clCreateKernel(prog, "kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4", &err), err));
         CL_CHECK(clReleaseProgram(prog));
         GGML_LOG_CONT(".");
     }
@@ -9646,6 +9691,8 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
             return op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]);
         case GGML_OP_IM2COL:
             return true;
+        case GGML_OP_IM2COL_3D:
+            return true;
         case GGML_OP_ARGSORT: {
             load_cl_kernels_argsort(backend_ctx);
 
@@ -20430,11 +20477,13 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
             return;
         }
 
-        // dp4a (int8) dense prefill GEMM, default off
+        static const char * q4_0_act_vload4_env = getenv("GGML_OPENCL_Q4_0_DP4A_ACT_VLOAD4");
+        const bool q4_0_act_vload4_on = q4_0_act_vload4_env == nullptr || atoi(q4_0_act_vload4_env) != 0;
+        // Enable dense prefill by default on X2E when the uint4 activation tile is enabled.
         static const char * q4_0_dense_dp4a_env = getenv("GGML_OPENCL_Q4_0_DENSE_DP4A");
         bool q4_0_dense_dp4a_on = q4_0_dense_dp4a_env
             ? (atoi(q4_0_dense_dp4a_env) != 0)
-            : false;
+            : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E && q4_0_act_vload4_on);
         // dot prod has to be available
         q4_0_dense_dp4a_on = backend_ctx->has_integer_dot && q4_0_dense_dp4a_on;
 
@@ -20461,7 +20510,9 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
             size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
             backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
 
-            cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a;
+            cl_kernel dk = q4_0_act_vload4_on
+                         ? backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4
+                         : backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a;
             int ai = 0;
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q4_0->q));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q4_0->d));
@@ -21754,8 +21805,13 @@ static void ggml_cl_mul_mat_q8_0_f32_adreno(ggml_backend_t backend, const ggml_t
                 }
             }
 
-            cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg
-                                    : backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a;
+            static const char * q8_0_act_vload4_env = getenv("GGML_OPENCL_Q8_0_DP4A_ACT_VLOAD4");
+            const bool q8_0_act_vload4_on = q8_0_act_vload4_env == nullptr || atoi(q8_0_act_vload4_env) != 0;
+            cl_kernel dk = use_wimg
+                ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg
+                : q8_0_act_vload4_on
+                    ? backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4
+                    : backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a;
             int ai = 0;
             if (use_wimg) {
                 CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q8_q_img));
@@ -22703,11 +22759,16 @@ static void ggml_cl_mul_mat_q4_k_f32_adreno(ggml_backend_t backend, const ggml_t
 
             const bool use_narrow = backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow != nullptr
                                  && N <= backend_ctx->max_n_dp4a_narrow_q4_k;
+            static const char * q4_k_act_vload4_env = getenv("GGML_OPENCL_Q4_K_DP4A_ACT_VLOAD4");
+            const bool q4_k_act_vload4_on = q4_k_act_vload4_env == nullptr || atoi(q4_k_act_vload4_env) != 0;
             cl_kernel dk = use_narrow
                 ? (use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg
                             : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow)
-                : (use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
-                            : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a);
+                : use_wimg
+                    ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
+                    : q4_k_act_vload4_on
+                        ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4
+                        : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a;
             int ai = 0;
             if (use_wimg) {
                 CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q4k_q_img));
@@ -23369,7 +23430,11 @@ static void ggml_cl_mul_mat_q6_K_f32_adreno(ggml_backend_t backend, const ggml_t
             size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
             backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
 
-            cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a;
+            static const char * q6_k_act_vload4_env = getenv("GGML_OPENCL_Q6_K_DP4A_ACT_VLOAD4");
+            const bool q6_k_act_vload4_on = q6_k_act_vload4_env == nullptr || atoi(q6_k_act_vload4_env) != 0;
+            cl_kernel dk = q6_k_act_vload4_on
+                         ? backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4
+                         : backend_ctx->kernel_gemm_noshuffle_q6_k_q8_1_dp4a;
             int ai = 0;
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q6_K->ql));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q6_K->qh));
@@ -23925,7 +23990,11 @@ static void ggml_cl_mul_mat_q5_K_f32_adreno(ggml_backend_t backend, const ggml_t
             size_t q_global[1] = { (size_t)(((n_blocks + 63) / 64) * 64) };
             backend_ctx->enqueue_ndrange_kernel(qk, 1, q_global, q_local, dst);
 
-            cl_kernel dk = backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a;
+            static const char * q5_k_act_vload4_env = getenv("GGML_OPENCL_Q5_K_DP4A_ACT_VLOAD4");
+            const bool q5_k_act_vload4_on = q5_k_act_vload4_env == nullptr || atoi(q5_k_act_vload4_env) != 0;
+            cl_kernel dk = q5_k_act_vload4_on
+                         ? backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4
+                         : backend_ctx->kernel_gemm_noshuffle_q5_k_q8_1_dp4a;
             int ai = 0;
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q5_k->q));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem),   &extra0_q5_k->qh));
@@ -30000,6 +30069,101 @@ static void ggml_cl_im2col(ggml_backend_t backend, const ggml_tensor * src0, con
     backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 }
 
+static void ggml_cl_im2col_3d(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_ASSERT(src0);
+    GGML_ASSERT(src1);
+    GGML_ASSERT(src1->extra);
+    GGML_ASSERT(dst);
+    GGML_ASSERT(dst->extra);
+
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_F32);
+
+    ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+    ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra;
+    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
+
+    cl_ulong offset1 = extra1->offset + src1->view_offs;
+    cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+    const int32_t s0 = ((const int32_t*)(dst->op_params))[0];
+    const int32_t s1 = ((const int32_t*)(dst->op_params))[1];
+    const int32_t s2 = ((const int32_t*)(dst->op_params))[2];
+    const int32_t p0 = ((const int32_t*)(dst->op_params))[3];
+    const int32_t p1 = ((const int32_t*)(dst->op_params))[4];
+    const int32_t p2 = ((const int32_t*)(dst->op_params))[5];
+    const int32_t d0 = ((const int32_t*)(dst->op_params))[6];
+    const int32_t d1 = ((const int32_t*)(dst->op_params))[7];
+    const int32_t d2 = ((const int32_t*)(dst->op_params))[8];
+    const cl_long IC = ((const int32_t*)(dst->op_params))[9];
+
+    const cl_long N  = src1->ne[3] / IC;
+    const cl_long ID = src1->ne[2];
+    const cl_long IH = src1->ne[1];
+    const cl_long IW = src1->ne[0];
+
+    const cl_long KD = src0->ne[2];
+    const cl_long KH = src0->ne[1];
+    const cl_long KW = src0->ne[0];
+
+    const cl_long OD = dst->ne[3] / N;
+    const cl_long OH = dst->ne[2];
+    const cl_long OW = dst->ne[1];
+
+    const size_t  es       = ggml_element_size(src1);
+    const cl_long stride_x = src1->nb[0] / es;
+    const cl_long stride_y = src1->nb[1] / es;
+    const cl_long stride_z = src1->nb[2] / es;
+    const cl_long stride_q = src1->nb[3] / es;
+
+    const cl_long IC_KD_KH_KW = IC * KD * KH * KW;
+
+    cl_kernel kernel;
+
+    if (dst->type == GGML_TYPE_F16) {
+        kernel = backend_ctx->kernel_im2col_3d_f16;
+    } else {
+        kernel = backend_ctx->kernel_im2col_3d_f32;
+    }
+
+    CL_CHECK(clSetKernelArg(kernel,   0, sizeof(cl_mem),   &extra1->data_device));
+    CL_CHECK(clSetKernelArg(kernel,   1, sizeof(cl_ulong), &offset1));
+    CL_CHECK(clSetKernelArg(kernel,   2, sizeof(cl_mem),   &extrad->data_device));
+    CL_CHECK(clSetKernelArg(kernel,   3, sizeof(cl_ulong), &offsetd));
+    CL_CHECK(clSetKernelArg(kernel,   4, sizeof(cl_long),  &IC));
+    CL_CHECK(clSetKernelArg(kernel,   5, sizeof(cl_long),  &ID));
+    CL_CHECK(clSetKernelArg(kernel,   6, sizeof(cl_long),  &IH));
+    CL_CHECK(clSetKernelArg(kernel,   7, sizeof(cl_long),  &IW));
+    CL_CHECK(clSetKernelArg(kernel,   8, sizeof(cl_long),  &KD));
+    CL_CHECK(clSetKernelArg(kernel,   9, sizeof(cl_long),  &KH));
+    CL_CHECK(clSetKernelArg(kernel,  10, sizeof(cl_long),  &KW));
+    CL_CHECK(clSetKernelArg(kernel,  11, sizeof(cl_long),  &OD));
+    CL_CHECK(clSetKernelArg(kernel,  12, sizeof(cl_long),  &OH));
+    CL_CHECK(clSetKernelArg(kernel,  13, sizeof(cl_long),  &OW));
+    CL_CHECK(clSetKernelArg(kernel,  14, sizeof(cl_long),  &N));
+    CL_CHECK(clSetKernelArg(kernel,  15, sizeof(cl_long),  &stride_x));
+    CL_CHECK(clSetKernelArg(kernel,  16, sizeof(cl_long),  &stride_y));
+    CL_CHECK(clSetKernelArg(kernel,  17, sizeof(cl_long),  &stride_z));
+    CL_CHECK(clSetKernelArg(kernel,  18, sizeof(cl_long),  &stride_q));
+    CL_CHECK(clSetKernelArg(kernel,  19, sizeof(int),      &s0));
+    CL_CHECK(clSetKernelArg(kernel,  20, sizeof(int),      &s1));
+    CL_CHECK(clSetKernelArg(kernel,  21, sizeof(int),      &s2));
+    CL_CHECK(clSetKernelArg(kernel,  22, sizeof(int),      &p0));
+    CL_CHECK(clSetKernelArg(kernel,  23, sizeof(int),      &p1));
+    CL_CHECK(clSetKernelArg(kernel,  24, sizeof(int),      &p2));
+    CL_CHECK(clSetKernelArg(kernel,  25, sizeof(int),      &d0));
+    CL_CHECK(clSetKernelArg(kernel,  26, sizeof(int),      &d1));
+    CL_CHECK(clSetKernelArg(kernel,  27, sizeof(int),      &d2));
+
+    const size_t local_col = 64;
+    const size_t num_blocks = (IC_KD_KH_KW + local_col - 1) / local_col;
+    size_t global_work_size[] = { num_blocks * local_col, (size_t)OW, (size_t)(N * OD * OH) };
+    size_t local_work_size[] = { local_col, 1, 1 };
+
+    backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+}
+
 static void ggml_cl_argsort(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -30901,6 +31065,12 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
             }
             func = ggml_cl_im2col;
             break;
+        case GGML_OP_IM2COL_3D:
+            if (!any_on_device) {
+                return false;
+            }
+            func = ggml_cl_im2col_3d;
+            break;
         case GGML_OP_ARGSORT:
             if (!any_on_device) {
                 return false;

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q4_0_q8_1_dp4a.cl

@@ -125,3 +125,129 @@ kernel void kernel_gemm_noshuffle_q4_0_q8_1_dp4a(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q4_0_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_q,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global const half   * src1_sa,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+    __local half sh_s[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const float d_w = (float)src0_d[rrow + sub * (uint)m];
+
+        const uint qsbase = rrow + (step >> 2) * (uint)m;
+        uint8 qw;
+        qw.s0 = EXP4(src0_q[qsbase + 0 * m]);
+        qw.s1 = EXP4(src0_q[qsbase + 1 * m]);
+        qw.s2 = EXP4(src0_q[qsbase + 2 * m]);
+        qw.s3 = EXP4(src0_q[qsbase + 3 * m]);
+        qw.s4 = EXP4(src0_q[qsbase + 4 * m]);
+        qw.s5 = EXP4(src0_q[qsbase + 5 * m]);
+        qw.s6 = EXP4(src0_q[qsbase + 6 * m]);
+        qw.s7 = EXP4(src0_q[qsbase + 7 * m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+
+            acc[g] += d_w * (LD4(sh_d, b) * rf - 8.0f * LD4(sh_s, b));
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q4_k_q8_1_dp4a.cl

@@ -279,3 +279,141 @@ kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_q,
+        __global const uchar  * src0_s,
+        __global const half   * src0_d,
+        __global const half   * src0_dm,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global const half   * src1_sa,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k,
+        uchar  mask_d6,
+        uchar  mask_d4,
+        uchar  mask_hi2
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+    __local half sh_s[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub     = step >> 5;
+        const uint sb_idx  = step / QK_K;
+        const uint sub_idx = sub & 7;
+
+        const float dd  = (float)src0_d [rrow + sb_idx * m];
+        const float dmm = (float)src0_dm[rrow + sb_idx * m];
+        global const uchar * sc = src0_s + sb_idx * K_SCALE_SIZE * (uint)m + rrow;
+        uchar sv, mn;
+        get_scale_min_k4(sub_idx, sc, (uint)m, &sv, &mn, mask_d6, mask_d4, mask_hi2);
+        const float scale = dd  * (float)sv;
+        const float minv  = dmm * (float)mn;
+
+        const uint wbase = rrow + (step >> 2) * (uint)m;
+        uint8 qw;
+        qw.s0 = EXP4(src0_q[wbase + 0 * m]);
+        qw.s1 = EXP4(src0_q[wbase + 1 * m]);
+        qw.s2 = EXP4(src0_q[wbase + 2 * m]);
+        qw.s3 = EXP4(src0_q[wbase + 3 * m]);
+        qw.s4 = EXP4(src0_q[wbase + 4 * m]);
+        qw.s5 = EXP4(src0_q[wbase + 5 * m]);
+        qw.s6 = EXP4(src0_q[wbase + 6 * m]);
+        qw.s7 = EXP4(src0_q[wbase + 7 * m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q5_k_q8_1_dp4a.cl

@@ -162,3 +162,141 @@ kernel void kernel_gemm_noshuffle_q5_k_q8_1_dp4a(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q5_k_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_q,
+        __global const uchar  * src0_qh,
+        __global const uchar  * src0_s,
+        __global const half   * src0_d,
+        __global const half   * src0_dm,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global const half   * src1_sa,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k,
+        uchar  mask_d6,
+        uchar  mask_d4,
+        uchar  mask_hi2
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint num_superblocks = (uint)k / QK_K;
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+    __local half sh_s[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub     = step >> 5;
+        const uint sb_idx  = step / QK_K;
+        const uint sub_idx = sub & 7;
+
+        const float dd  = (float)src0_d [rrow + sb_idx * m];
+        const float dmm = (float)src0_dm[rrow + sb_idx * m];
+        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
+        uchar sv, mn;
+        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
+        const float scale = dd  * (float)sv;
+        const float minv  = dmm * (float)mn;
+
+        const uint wbase  = rrow + (step >> 2) * (uint)m;
+        const uint qhbase = rrow + (step >> 3) * (uint)m;
+        uint8 qw;
+#define QWU(u) ( EXP4((uint)src0_q[wbase + (uint)(u) * m]) \
+               | EXP1( (uint)((src0_qh[qhbase + (uint)((u) >> 1) * m] >> (((u) & 1) * 4)) & 0x0Fu) ) )
+        qw.s0 = QWU(0); qw.s1 = QWU(1); qw.s2 = QWU(2); qw.s3 = QWU(3);
+        qw.s4 = QWU(4); qw.s5 = QWU(5); qw.s6 = QWU(6); qw.s7 = QWU(7);
+#undef QWU
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+            sh_s[lid] = (c < (uint)n_no_padding) ? src1_sa[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+            acc[g] += scale * LD4(sh_d, b) * rf - minv * LD4(sh_s, b);
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q6_k_q8_1_dp4a.cl

@@ -142,3 +142,117 @@ kernel void kernel_gemm_noshuffle_q6_k_q8_1_dp4a(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q6_k_q8_1_dp4a_act_vload4(
+        __global const ushort * src0_ql,
+        __global const uchar  * src0_qh,
+        __global const ushort * src0_s,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub    = step >> 5;
+        const uint sb_idx = step / QK_K;
+
+        const float dd = (float)src0_d[rrow + sb_idx * m];
+        const char2 sc = as_char2(src0_s[rrow + sub * m]);
+        const float scale0 = dd * (float)sc.s0;
+        const float scale1 = dd * (float)sc.s1;
+
+        const uint wbase = rrow + (step >> 2) * (uint)m;
+        uint qw[8];
+        #pragma unroll
+        for (int u = 0; u < 8; ++u) {
+            const uint o  = wbase + (uint)u * (uint)m;
+            qw[u] = SIGN6(EXP4((uint)src0_ql[o]) | EXP2((uint)src0_qh[o]));
+        }
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            #define DOT_TOK(j) { \
+                const uint4 a0 = sh_qa4[b + (j)][0]; \
+                const uint4 a1 = sh_qa4[b + (j)][1]; \
+                const int raw1 = dot4_q8a(qw[0], qw[1], qw[2], qw[3], a0.x, a0.y, a0.z, a0.w); \
+                const int raw2 = dot4_q8a(qw[4], qw[5], qw[6], qw[7], a1.x, a1.y, a1.z, a1.w); \
+                rf.s##j = scale0 * (float)raw1 + scale1 * (float)raw2; \
+            }
+            DOT_TOK(0); DOT_TOK(1); DOT_TOK(2); DOT_TOK(3);
+            #undef DOT_TOK
+            const float4 ad = (float4)((float)sh_d[b+0], (float)sh_d[b+1], (float)sh_d[b+2], (float)sh_d[b+3]);
+            acc[g] += ad * rf;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/gemm_noshuffle_q8_0_q8_1_dp4a.cl

@@ -210,3 +210,125 @@ kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a_wimg(
     }
 #undef NGROUPS
 }
+
+// Stage each 32-K activation block as two uint4 values.
+inline int dot8_q8a_v(uint8 qw, uint4 a0, uint4 a1) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a0.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a0.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a0.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a0.w, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s4, a1.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s5, a1.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s6, a1.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s7, a1.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q8_0_q8_1_dp4a_act_vload4(
+        __global const uint  * src0_q,
+        __global const half  * src0_d,
+        __global const uint  * src1_qa,
+        __global const half  * src1_da,
+        __global       float * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        acc[g] = (float4)(0.0f);
+    }
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+
+        const float d_w = (float)src0_d[rrow + sub * (uint)m];
+
+        const uint wbase = rrow + (step >> 2) * (uint)m;
+        uint8 qw;
+        qw.s0 = src0_q[wbase + 0 * m];
+        qw.s1 = src0_q[wbase + 1 * m];
+        qw.s2 = src0_q[wbase + 2 * m];
+        qw.s3 = src0_q[wbase + 3 * m];
+        qw.s4 = src0_q[wbase + 4 * m];
+        qw.s5 = src0_q[wbase + 5 * m];
+        qw.s6 = src0_q[wbase + 6 * m];
+        qw.s7 = src0_q[wbase + 7 * m];
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define DOTV(T) dot8_q8a_v(qw, sh_qa4[T][0], sh_qa4[T][1])
+            rf.s0 = (float)DOTV(b+0);
+            rf.s1 = (float)DOTV(b+1);
+            rf.s2 = (float)DOTV(b+2);
+            rf.s3 = (float)DOTV(b+3);
+#undef DOTV
+            acc[g] += d_w * LD4(sh_d, b) * rf;
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) {
+            dst[(c0 + 0) * (uint)m + row] = a.s0;
+        }
+        if (c0 + 1 < (uint)n_no_padding) {
+            dst[(c0 + 1) * (uint)m + row] = a.s1;
+        }
+        if (c0 + 2 < (uint)n_no_padding) {
+            dst[(c0 + 2) * (uint)m + row] = a.s2;
+        }
+        if (c0 + 3 < (uint)n_no_padding) {
+            dst[(c0 + 3) * (uint)m + row] = a.s3;
+        }
+    }
+#undef NGROUPS
+}

ggml/src/ggml-opencl/kernels/im2col_3d_f16.cl

@@ -0,0 +1,90 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+kernel void kernel_im2col_3d_f16(
+        global float * src1,
+        ulong offset1,
+        global half * dst,
+        ulong offsetd,
+        long IC,
+        long ID,
+        long IH,
+        long IW,
+        long KD,
+        long KH,
+        long KW,
+        long OD,
+        long OH,
+        long OW,
+        long N,
+        long stride_x,
+        long stride_y,
+        long stride_z,
+        long stride_q,
+        int s0,
+        int s1,
+        int s2,
+        int p0,
+        int p1,
+        int p2,
+        int d0,
+        int d1,
+        int d2
+) {
+    const int ic = (int)IC;
+    const int kd = (int)KD;
+    const int kh = (int)KH;
+    const int kw = (int)KW;
+    const int od = (int)OD;
+    const int oh = (int)OH;
+    const int ow = (int)OW;
+    const int n  = (int)N;
+
+    const int KD_KH_KW    = kd * kh * kw;
+    const int KH_KW       = kh * kw;
+    const int IC_KD_KH_KW = ic * KD_KH_KW;
+
+    const int i = (int)get_global_id(0);
+    if (i >= IC_KD_KH_KW) {
+        return;
+    }
+
+    const int iow = (int)get_global_id(1);
+    if (iow >= ow) {
+        return;
+    }
+
+    const int OD_OH = od * oh;
+    const int iz = (int)get_global_id(2);
+    if (iz >= n * OD_OH) {
+        return;
+    }
+
+    src1 = (global float*)((global char*)src1 + offset1);
+    dst  = (global half*)((global char*)dst + offsetd);
+
+    const int iic = i / KD_KH_KW;
+    const int ikd = (i - iic * KD_KH_KW) / KH_KW;
+    const int ikh = (i - iic * KD_KH_KW - ikd * KH_KW) / kw;
+    const int ikw = i % kw;
+
+    const int in  = iz / OD_OH;
+    const int iod = (iz - in * OD_OH) / oh;
+    const int ioh = iz % oh;
+
+    const int iiw = iow * s0 + ikw * d0 - p0;
+    const int iih = ioh * s1 + ikh * d1 - p1;
+    const int iid = iod * s2 + ikd * d2 - p2;
+
+    const long offset_dst =
+        (long)(iz * ow + iow) * (long)IC_KD_KH_KW +
+        (long)(iic * KD_KH_KW + ikd * KH_KW + ikh * kw + ikw);
+
+    if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW || iid < 0 || iid >= ID) {
+        dst[offset_dst] = 0.0f;
+    } else {
+        const long offset_src =
+            (long)(in * ic + iic) * stride_q + (long)iid * stride_z +
+            (long)iih * stride_y + (long)iiw * stride_x;
+        dst[offset_dst] = convert_half_rte(src1[offset_src]);
+    }
+}

ggml/src/ggml-opencl/kernels/im2col_3d_f32.cl

@@ -0,0 +1,90 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+kernel void kernel_im2col_3d_f32(
+        global float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        long IC,
+        long ID,
+        long IH,
+        long IW,
+        long KD,
+        long KH,
+        long KW,
+        long OD,
+        long OH,
+        long OW,
+        long N,
+        long stride_x,
+        long stride_y,
+        long stride_z,
+        long stride_q,
+        int s0,
+        int s1,
+        int s2,
+        int p0,
+        int p1,
+        int p2,
+        int d0,
+        int d1,
+        int d2
+) {
+    const int ic = (int)IC;
+    const int kd = (int)KD;
+    const int kh = (int)KH;
+    const int kw = (int)KW;
+    const int od = (int)OD;
+    const int oh = (int)OH;
+    const int ow = (int)OW;
+    const int n  = (int)N;
+
+    const int KD_KH_KW    = kd * kh * kw;
+    const int KH_KW       = kh * kw;
+    const int IC_KD_KH_KW = ic * KD_KH_KW;
+
+    const int i = (int)get_global_id(0);
+    if (i >= IC_KD_KH_KW) {
+        return;
+    }
+
+    const int iow = (int)get_global_id(1);
+    if (iow >= ow) {
+        return;
+    }
+
+    const int OD_OH = od * oh;
+    const int iz = (int)get_global_id(2);
+    if (iz >= n * OD_OH) {
+        return;
+    }
+
+    src1 = (global float*)((global char*)src1 + offset1);
+    dst  = (global float*)((global char*)dst + offsetd);
+
+    const int iic = i / KD_KH_KW;
+    const int ikd = (i - iic * KD_KH_KW) / KH_KW;
+    const int ikh = (i - iic * KD_KH_KW - ikd * KH_KW) / kw;
+    const int ikw = i % kw;
+
+    const int in  = iz / OD_OH;
+    const int iod = (iz - in * OD_OH) / oh;
+    const int ioh = iz % oh;
+
+    const int iiw = iow * s0 + ikw * d0 - p0;
+    const int iih = ioh * s1 + ikh * d1 - p1;
+    const int iid = iod * s2 + ikd * d2 - p2;
+
+    const long offset_dst =
+        (long)(iz * ow + iow) * (long)IC_KD_KH_KW +
+        (long)(iic * KD_KH_KW + ikd * KH_KW + ikh * kw + ikw);
+
+    if (iih < 0 || iih >= IH || iiw < 0 || iiw >= IW || iid < 0 || iid >= ID) {
+        dst[offset_dst] = 0.0f;
+    } else {
+        const long offset_src =
+            (long)(in * ic + iic) * stride_q + (long)iid * stride_z +
+            (long)iih * stride_y + (long)iiw * stride_x;
+        dst[offset_dst] = src1[offset_src];
+    }
+}

ggml/src/ggml-sycl/convert.cpp

@@ -177,15 +177,12 @@ static void dequantize_row_q8_0_sycl_reorder(const void *vx, dst_t *y, const int
     dpct::has_capability_or_fail(stream->get_device(),
                                     {sycl::aspect::fp16});
 
+    // flat range<1>, two pairs (4 elements) per item: no nd_range/group math, one vectorized store.
+    // k % QK8_0 == 0 implies k % 4 == 0, so the wide launch covers the row exactly (no tail).
     GGML_ASSERT(k % QK8_0 == 0);
-    const int64_t n_groups = ceil_div(k / 2, SYCL_DEQUANTIZE_BLOCK_SIZE);
-    // decrease global range when it exceeds the max int; the kernel loops to cover the rest
-    const int64_t wg_size  = downsample_sycl_global_range(n_groups, SYCL_DEQUANTIZE_BLOCK_SIZE);
-    stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, n_groups) *
-        sycl::range<3>(1, 1, wg_size),
-        sycl::range<3>(1, 1, wg_size)),
-        [=](sycl::nd_item<3> item_ct1) {
-            dequantize_block_q8_0_reorder(vx, y, k, item_ct1);
+    stream->parallel_for(sycl::range<1>((size_t) (k / 4)),
+        [=](sycl::id<1> idx) {
+            dequantize_block_q8_0_reorder_wide(vx, y, k, idx);
         });
 }
 

ggml/src/ggml-sycl/dequantize.hpp

@@ -857,6 +857,26 @@ static void dequantize_block_q8_0_reorder(const void * __restrict__ vx, dst_t *
     }
 }
 
+template<typename dst_t>
+static void dequantize_block_q8_0_reorder_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t k,
+                                               const sycl::id<1> idx) {
+    const int64_t i = 4 * (int64_t) idx[0];
+
+    const int8_t *     qs    = (const int8_t *) vx;
+    const sycl::half * d_ptr = (const sycl::half *) ((const uint8_t *) vx + k);
+
+    const float d0 = float(d_ptr[i / QK8_0]);
+    const float d1 = float(d_ptr[(i + 2) / QK8_0]);
+
+    sycl::vec<float, 4> p;
+    p[0] = d0 * (float) qs[i + 0];
+    p[1] = d0 * (float) qs[i + 1];
+    p[2] = d1 * (float) qs[i + 2];
+    p[3] = d1 * (float) qs[i + 3];
+
+    *reinterpret_cast<sycl::vec<dst_t, 4> *>(yy + i) = p.convert<dst_t, sycl::rounding_mode::automatic>();
+}
+
 template<typename dst_t>
 static void dequantize_block_q4_1(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb32,
                                   const sycl::nd_item<3> &item_ct1) {

ggml/src/ggml-sycl/fattn.cpp

@@ -281,6 +281,22 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
 void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     ggml_sycl_set_device(ctx.device);
 
+    ggml_tensor * Q = dst->src[0];
+    ggml_tensor * K = dst->src[1];
+    ggml_tensor * V = dst->src[2];
+    if (K && Q && K->ne[0] != Q->ne[0] && Q->ne[0] > 0 && K->ne[0] % Q->ne[0] == 0) {
+        const int64_t d = Q->ne[0];
+        const int64_t n_heads = (K->ne[0] * K->ne[2]) / d;
+        K->ne[0] = d;
+        K->ne[2] = n_heads;
+    }
+    if (V && Q && V->ne[0] != Q->ne[0] && Q->ne[0] > 0 && V->ne[0] % Q->ne[0] == 0) {
+        const int64_t d = Q->ne[0];
+        const int64_t n_heads = (V->ne[0] * V->ne[2]) / d;
+        V->ne[0] = d;
+        V->ne[2] = n_heads;
+    }
+
     // sparse nodes are gathered down to n_kv_max rows and re-dispatched here
     if (ggml_sycl_flash_attn_ext_sparse(ctx, dst)) {
         return;

ggml/src/ggml.c

@@ -1095,6 +1095,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "DSV4_HC_COMB",
     "DSV4_HC_PRE",
     "DSV4_HC_POST",
+    "TOP_K_MASK",
 
     "UNARY",
 
@@ -1112,7 +1113,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GLU",
 };
 
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1210,6 +1211,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "dsv4_hc_comb(mixes, scale, base)",
     "dsv4_hc_pre(x, weights)",
     "dsv4_hc_post(x, residual, post, comb)",
+    "top_k_mask(mask, top_k, score, blocks, extra)",
 
     "unary(x)",
 
@@ -1227,7 +1229,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "glu(x)",
 };
 
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -6630,6 +6632,48 @@ struct ggml_tensor * ggml_dsv4_hc_post(
     return result;
 }
 
+// ggml_top_k_mask
+
+struct ggml_tensor * ggml_top_k_mask(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * mask,
+        struct ggml_tensor  * top_k,
+        struct ggml_tensor  * score,
+        struct ggml_tensor  * blocks,
+        struct ggml_tensor  * extra) {
+    GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32);
+    GGML_ASSERT(mask->nb[0] == ggml_type_size(mask->type));
+    GGML_ASSERT(top_k->type == GGML_TYPE_I32);
+    GGML_ASSERT(top_k->ne[1] % mask->ne[1] == 0 && top_k->ne[2] % mask->ne[2] == 0 && top_k->ne[3] % mask->ne[3] == 0);
+
+    if (score) {
+        GGML_ASSERT(score->type == GGML_TYPE_F32);
+        GGML_ASSERT(score->ne[1] == top_k->ne[1] && score->ne[2] == top_k->ne[2] && score->ne[3] == top_k->ne[3]);
+    }
+
+    if (blocks) {
+        GGML_ASSERT(blocks->type == GGML_TYPE_I32);
+        GGML_ASSERT(blocks->ne[2] == 1 && blocks->ne[3] == 1);
+        GGML_ASSERT(!score || blocks->ne[1] == score->ne[0]);
+    }
+
+    if (extra) {
+        GGML_ASSERT(extra->type == GGML_TYPE_I32);
+        GGML_ASSERT(extra->ne[1] == top_k->ne[1] && extra->ne[2] == top_k->ne[2] && extra->ne[3] == top_k->ne[3]);
+    }
+
+    struct ggml_tensor * result = ggml_new_tensor_4d(ctx, mask->type, mask->ne[0], top_k->ne[1], top_k->ne[2], top_k->ne[3]);
+
+    result->op     = GGML_OP_TOP_K_MASK;
+    result->src[0] = mask;
+    result->src[1] = top_k;
+    result->src[2] = score;
+    result->src[3] = blocks;
+    result->src[4] = extra;
+
+    return result;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 struct ggml_hash_set ggml_hash_set_new(size_t size) {

src/llama-adapter.cpp

@@ -179,13 +179,13 @@ static void llama_adapter_lora_init_impl(llama_model & model, FILE * file, llama
                 ? format("%s[%s,%zu]", gguf_type_name(type), gguf_type_name(gguf_get_arr_type(gguf_ctx, i)), gguf_get_arr_n(gguf_ctx, i))
                 : gguf_type_name(type);
             const char * name = gguf_get_key(gguf_ctx, i);
-            const std::string value = gguf_kv_to_str(gguf_ctx, i);
+            const size_t MAX_VALUE_LEN = 40;
+            const std::string value = gguf_kv_to_str(gguf_ctx, i, MAX_VALUE_LEN);
 
             if (type != GGUF_TYPE_ARRAY) {
                 adapter.gguf_kv.emplace(name, value);
             }
 
-            const size_t MAX_VALUE_LEN = 40;
             std::string print_value = value.size() > MAX_VALUE_LEN ? format("%s...", value.substr(0, MAX_VALUE_LEN - 3).c_str()) : value;
             replace_all(print_value, "\n", "\\n");
 

src/llama-arch.cpp

@@ -1110,7 +1110,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_MASKED_EMBD_ORDERING,       {LLM_TENSOR_LAYER_INPUT,     GGML_OP_NONE}},
     // eagle3
     {LLM_TENSOR_FC,                         {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
-    {LLM_TENSOR_D2T,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
+    {LLM_TENSOR_D2T,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_SET_ROWS}},
     // dspark
     {LLM_TENSOR_DSPARK_MARKOV_W1,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
     {LLM_TENSOR_DSPARK_MARKOV_W2,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},

src/llama-context.cpp

@@ -165,6 +165,12 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = {
     /*.n_tokens_per_seq =*/ 1,
 };
 
+static const llm_fused_op_probe llm_fused_op_top_k_mask_probe = {
+    /*.op               =*/ LLM_FUSED_OP_TOP_K_MASK,
+    /*.name             =*/ "fused top-k mask",
+    /*.n_tokens_per_seq =*/ 1,
+};
+
 llama_context::llama_context(
         const llama_model & model,
               llama_context_params params) :
@@ -328,6 +334,10 @@ llama_context::llama_context(
     cparams.fused_dsv4_hc_post = true;
     cparams.auto_fhc           = true;
 
+    // only the indexer models build this mask, skip the probe for the others
+    cparams.fused_top_k_mask = true;
+    cparams.auto_ftkm        = hparams.indexer_top_k > 0;
+
     // with causal attention, the batch size is limited by the context size
     cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
 
@@ -671,6 +681,12 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
         resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post);
         cparams.auto_fhc = false;
     }
+
+    if (cparams.auto_ftkm) {
+        LLAMA_LOG_INFO("%s: resolving fused top-k mask support:\n", func);
+        resolve(llm_fused_op_top_k_mask_probe, cparams.fused_top_k_mask);
+        cparams.auto_ftkm = false;
+    }
 }
 
 static int llama_graph_n_input_tensors(ggml_cgraph * gf) {

src/llama-cparams.h

@@ -48,6 +48,8 @@ struct llama_cparams {
     bool fused_dsv4_hc_comb;
     bool fused_dsv4_hc_post;
     bool auto_fhc;
+    bool fused_top_k_mask;   // use fused top-k sparse attention mask
+    bool auto_ftkm;
     bool no_perf;
     bool warmup;             // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
     bool op_offload;

src/llama-graph.cpp

@@ -1666,6 +1666,40 @@ ggml_tensor * llm_graph_context::build_lora_mm(
     return res;
 }
 
+llm_graph_amax_src1 llm_graph_context::build_amax_src1_f16(
+          ggml_tensor * cur) const {
+    llm_graph_amax_src1 result = {
+        /*.res     =*/ cur,
+        /*.scale_s =*/ nullptr,
+    };
+
+    const int64_t n_embd = cur->ne[0];
+
+    // amax = max|cur| along the rows for each [expert, token] -> [1, n_expert, n_tokens]
+    ggml_tensor * ab = ggml_abs(ctx0, cur);
+    ggml_tensor * amax = ggml_pool_1d(ctx0, ab, GGML_OP_POOL_MAX, n_embd, n_embd, 0);
+
+    // ref: https://github.com/ggml-org/llama.cpp/pull/26223
+    constexpr float k_log2_e = 1.4426950408889634f; // 1 / ln(2)
+    constexpr float k_ln_2   = 0.6931471805599453f;
+
+    ggml_tensor * e;
+    e = ggml_scale(ctx0, amax, 1.0f/32768.0f);
+    e = ggml_log  (ctx0, e);
+    e = ggml_scale(ctx0, e, k_log2_e);
+    e = ggml_ceil (ctx0, e);
+    e = ggml_clamp(ctx0, e, 0.0f, 127.0f);
+    e = ggml_scale(ctx0, e, k_ln_2);
+    e = ggml_exp  (ctx0, e);
+
+    result = {
+        /*.res     =*/ ggml_div(ctx0, cur, e),
+        /*.scale_s =*/ e,
+    };
+
+    return result;
+}
+
 ggml_tensor * llm_graph_context::build_lora_mm_id(
           ggml_tensor * w,   // ggml_tensor * as
           ggml_tensor * cur, // ggml_tensor * b
@@ -2465,12 +2499,25 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
             GGML_ABORT("fatal error");
     }
 
-    experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
-    if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
-        // src1 can exceed F16 range
-        ggml_prec_set_src(experts, GGML_PREC_F32, 1);
+    {
+        // TODO: this graph-level rescaling should only be applied when the corresponding backend does not support
+        //       ggml_prec_set_src(experts, GGML_PREC_BF16, 1);
+        //       ref: https://github.com/ggml-org/llama.cpp/pull/30656#issuecomment-6269597922
+        llm_graph_amax_src1 amax_src1 = { nullptr, nullptr };
+
+        if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
+            amax_src1 = build_amax_src1_f16(cur);
+            cur = amax_src1.res;
+        }
+
+        experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
+        //ggml_prec_set_src(experts, GGML_PREC_BF16, 1); // (see TODO above)
+        cb(experts, "ffn_moe_down", il);
+
+        if (amax_src1.scale_s) {
+            experts = ggml_mul(ctx0, experts, amax_src1.scale_s);
+        }
     }
-    cb(experts, "ffn_moe_down", il);
 
     if (down_exps_s) {
         cb(experts, "ffn_moe_down_scaled", il);
@@ -3328,6 +3375,27 @@ ggml_tensor * llm_graph_context::build_attn(
     return cur;
 }
 
+ggml_tensor * llm_graph_context::build_fused_top_k_mask(
+        ggml_tensor * kq_mask,
+        ggml_tensor * top_k,
+        ggml_tensor * score,
+        ggml_tensor * blocks,
+        ggml_tensor * extra,
+                int   il) const {
+    const int64_t m1 = kq_mask->ne[1];
+    const int64_t m2 = kq_mask->ne[2];
+    const int64_t m3 = kq_mask->ne[3];
+
+    auto rows = [&](ggml_tensor * t) {
+        return t ? ggml_reshape_4d(ctx0, t, t->ne[0], m1, m2, m3) : nullptr;
+    };
+
+    ggml_tensor * cur = ggml_top_k_mask(ctx0, kq_mask, rows(top_k), rows(score), blocks, rows(extra));
+    res->add_fused_node({LLM_FUSED_OP_TOP_K_MASK, cur, il});
+
+    return cur;
+}
+
 ggml_tensor * llm_graph_context::build_attn(
         llm_graph_input_attn_k_dsa * inp,
         ggml_tensor * wo,
@@ -3360,31 +3428,36 @@ ggml_tensor * llm_graph_context::build_attn(
 
     const auto & kq_mask = inp->get_kq_mask_mla();
 
-    // prepare new kq mask - starts filled with -INFINITY
-    ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
+    ggml_tensor * kq_mask_top_k = nullptr;
+    if (cparams.fused_top_k_mask) {
+        kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il);
+    } else {
+        // prepare new kq mask - starts filled with -INFINITY
+        ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
 
-    // reshape KQ mask into tensor with rows of size 1:
-    // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
-    kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
+        // reshape KQ mask into tensor with rows of size 1:
+        // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]
+        kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);
 
-    // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
-    ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
+        // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]
+        ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);
 
-    // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
-    // this will be our source of zero values for unmasking top k mask elements
-    ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
-    zeros = ggml_fill(ctx0, zeros, 0.0f);
+        // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]
+        // this will be our source of zero values for unmasking top k mask elements
+        ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);
+        zeros = ggml_fill(ctx0, zeros, 0.0f);
 
-    // modify KQ mask by unmasking elements that are in top_k indices
-    // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
-    ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
+        // modify KQ mask by unmasking elements that are in top_k indices
+        // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])
+        kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);
 
-    // reshape to restore the original shape of KQ mask:
-    // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
-    kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
+        // reshape to restore the original shape of KQ mask:
+        // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]
+        kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);
 
-    // combine with the original kq mask
-    kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+        // combine with the original kq mask
+        kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);
+    }
 
     ggml_tensor * q = q_cur;
     ggml_tensor * k = mctx_cur->get_k(ctx0, il);

src/llama-graph.h

@@ -53,6 +53,7 @@ enum llm_fused_op {
     LLM_FUSED_OP_DSV4_HC_PRE,
     LLM_FUSED_OP_DSV4_HC_COMB,
     LLM_FUSED_OP_DSV4_HC_POST,
+    LLM_FUSED_OP_TOP_K_MASK,
 };
 
 enum llm_ffn_op_type : int {
@@ -1012,6 +1013,11 @@ struct llm_graph_qkv {
     ggml_tensor * v; // [n_embd_head, n_head_kv, n_tokens]
 };
 
+struct llm_graph_amax_src1 {
+    ggml_tensor * res;     // cur rescaled
+    ggml_tensor * scale_s; // scale to undo on the matmul result, or nullptr
+};
+
 struct llm_graph_context {
     const llm_arch arch;
 
@@ -1115,6 +1121,11 @@ struct llm_graph_context {
               ggml_tensor * w_s   = nullptr,
               ggml_tensor * slots = nullptr) const;
 
+    // rescale `cur` (the src1 of a matmul) by a power of two, computed from its maximum absolute
+    // when the values already fit in F16, the factors are exactly 1.0f
+    llm_graph_amax_src1 build_amax_src1_f16(
+              ggml_tensor * cur) const;
+
     ggml_tensor * build_norm(
              ggml_tensor * cur,
              ggml_tensor * mw,
@@ -1325,6 +1336,15 @@ struct llm_graph_context {
                   float   kq_scale,
                     int   il) const;
 
+    // ggml_top_k_mask with top_k, score and extra viewed in the row layout of kq_mask, see ggml.h
+    ggml_tensor * build_fused_top_k_mask(
+            ggml_tensor * kq_mask, // [n_kv, n_tokens/n_stream, 1, n_stream]
+            ggml_tensor * top_k,   // [n_top, n_tokens]
+            ggml_tensor * score,   // [n_pick, n_tokens], can be null
+            ggml_tensor * blocks,  // [n_blk, n_pick], can be null
+            ggml_tensor * extra,   // [n_extra, n_tokens], can be null
+                    int   il) const;
+
     llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;
 
     llm_graph_input_dsv4 * build_inp_dsv4() const;

src/llama-impl.cpp

@@ -11,7 +11,6 @@
 #include <cstdarg>
 #include <cstring>
 #include <vector>
-#include <sstream>
 
 struct llama_logger_state {
     ggml_log_callback log_callback = llama_log_callback_default;
@@ -149,7 +148,7 @@ std::string llama_format_tensor_shape(const struct ggml_tensor * t) {
     return buf;
 }
 
-static std::string gguf_data_to_str(enum gguf_type type, const void * data, int i) {
+static std::string gguf_data_to_str(enum gguf_type type, const void * data, size_t i) {
     switch (type) {
         case GGUF_TYPE_UINT8:   return std::to_string(((const uint8_t  *)data)[i]);
         case GGUF_TYPE_INT8:    return std::to_string(((const int8_t   *)data)[i]);
@@ -166,7 +165,7 @@ static std::string gguf_data_to_str(enum gguf_type type, const void * data, int
     }
 }
 
-std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) {
+std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i, size_t max_len) {
     const enum gguf_type type = gguf_get_kv_type(ctx_gguf, i);
 
     switch (type) {
@@ -175,28 +174,38 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) {
         case GGUF_TYPE_ARRAY:
             {
                 const enum gguf_type arr_type = gguf_get_arr_type(ctx_gguf, i);
-                int arr_n = gguf_get_arr_n(ctx_gguf, i);
+                const size_t arr_n = gguf_get_arr_n(ctx_gguf, i);
                 const void * data = arr_type == GGUF_TYPE_STRING ? nullptr : gguf_get_arr_data(ctx_gguf, i);
-                std::stringstream ss;
-                ss << "[";
-                for (int j = 0; j < arr_n; j++) {
-                    if (arr_type == GGUF_TYPE_STRING) {
-                        std::string val = gguf_get_arr_str(ctx_gguf, i, j);
-                        // escape quotes
-                        replace_all(val, "\\", "\\\\");
-                        replace_all(val, "\"", "\\\"");
-                        ss << '"' << val << '"';
-                    } else if (arr_type == GGUF_TYPE_ARRAY) {
-                        ss << "???";
-                    } else {
-                        ss << gguf_data_to_str(arr_type, data, j);
+                std::string res = "[";
+                size_t j = 0;
+                for (; j < arr_n && res.size() <= max_len; j++) {
+                    if (j > 0) {
+                        res += ", ";
                     }
-                    if (j < arr_n - 1) {
-                        ss << ", ";
+                    if (arr_type == GGUF_TYPE_STRING) {
+                        const char * s = gguf_get_arr_str(ctx_gguf, i, j);
+                        const size_t need = res.size() + 2 * strlen(s) + 2;
+                        if (res.capacity() < need) {
+                            res.reserve(std::max(need, 2 * res.capacity()));
+                        }
+                        res += '"';
+                        for (; *s; s++) {
+                            if (*s == '\\' || *s == '"') {
+                                res += '\\';
+                            }
+                            res += *s;
+                        }
+                        res += '"';
+                    } else if (arr_type == GGUF_TYPE_ARRAY) {
+                        res += "???";
+                    } else {
+                        res += gguf_data_to_str(arr_type, data, j);
                     }
                 }
-                ss << "]";
-                return ss.str();
+                if (j == arr_n) {
+                    res += "]";
+                }
+                return res;
             }
         default:
             return gguf_data_to_str(type, gguf_get_val_data(ctx_gguf, i), 0);

src/llama-impl.h

@@ -107,4 +107,5 @@ std::string format(const char * fmt, ...);
 std::string llama_format_tensor_shape(const std::vector<int64_t> & ne);
 std::string llama_format_tensor_shape(const struct ggml_tensor * t);
 
-std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
+// arrays are rendered only until the result is longer than max_len
+std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i, size_t max_len = SIZE_MAX);

src/llama-model-loader.cpp

@@ -807,8 +807,8 @@ llama_model_loader::llama_model_loader(
                 ? format("%s[%s,%zu]", gguf_type_name(type), gguf_type_name(gguf_get_arr_type(metadata, i)), gguf_get_arr_n(metadata, i))
                 : gguf_type_name(type);
 
-            std::string value          = gguf_kv_to_str(metadata, i);
             const size_t MAX_VALUE_LEN = 40;
+            std::string value          = gguf_kv_to_str(metadata, i, MAX_VALUE_LEN);
             if (value.size() > MAX_VALUE_LEN) {
                 value = format("%s...", value.substr(0, MAX_VALUE_LEN - 3).c_str());
             }
@@ -951,6 +951,13 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
                 ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512);
                 op_tensor = ggml_get_rows(ctx, w, b);
             } break;
+        case GGML_OP_SET_ROWS:
+            {
+                // w holds the row indices (e.g. d2t)
+                ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, w->ne[0]);
+                ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, w->ne[0]);
+                op_tensor = ggml_set_rows(ctx, a, b, w);
+            } break;
         case GGML_OP_MUL_MAT:
             {
                 ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]);

src/llama-vocab.cpp

@@ -2260,21 +2260,22 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                     throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
                 }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
+                bpe_ranks.reserve(n_merges);
                 for (int i = 0; i < n_merges; i++) {
-                    const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
+                    const char * word = gguf_get_arr_str(ctx, merges_keyidx, i);
                     //GGML_ASSERT(unicode_cpts_from_utf8(word).size() > 0);
 
                     std::string first;
                     std::string second;
 
-                    const size_t pos = word.find(' ', 1);
+                    const char * sep = *word ? strchr(word + 1, ' ') : nullptr;
 
-                    if (pos != std::string::npos) {
-                        first  = word.substr(0, pos);
-                        second = word.substr(pos + 1);
+                    if (sep) {
+                        first.assign(word, sep - word);
+                        second.assign(sep + 1);
                     }
 
-                    bpe_ranks.emplace(std::make_pair(first, second), i);
+                    bpe_ranks.emplace(std::make_pair(std::move(first), std::move(second)), i);
                 }
             }
 
@@ -2371,20 +2372,21 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                     throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
                 }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
+                bpe_ranks.reserve(n_merges);
                 for (int i = 0; i < n_merges; i++) {
-                    const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
+                    const char * word = gguf_get_arr_str(ctx, merges_keyidx, i);
 
                     std::string first;
                     std::string second;
 
-                    const size_t pos = word.find(' ', 1);
+                    const char * sep = *word ? strchr(word + 1, ' ') : nullptr;
 
-                    if (pos != std::string::npos) {
-                        first  = word.substr(0, pos);
-                        second = word.substr(pos + 1);
+                    if (sep) {
+                        first.assign(word, sep - word);
+                        second.assign(sep + 1);
                     }
 
-                    bpe_ranks.emplace(std::make_pair(first, second), i);
+                    bpe_ranks.emplace(std::make_pair(std::move(first), std::move(second)), i);
                 }
             }
 
@@ -2764,6 +2766,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
     }
 
     id_to_token.resize(n_tokens);
+    token_to_id.reserve(n_tokens);
 
     for (uint32_t i = 0; i < n_tokens; i++) {
         std::string word = gguf_get_arr_str(ctx, token_idx, i);
@@ -3676,13 +3679,14 @@ static void llama_unescape_whitespace(std::string & word) {
 
 static std::string llama_decode_text(const std::string & text) {
     std::string decoded_text;
+    decoded_text.reserve(text.size());
 
     const auto cpts = unicode_cpts_from_utf8(text);
     for (const auto cpt : cpts) {
-        const auto utf8 = unicode_cpt_to_utf8(cpt);
         try {
-            decoded_text += unicode_utf8_to_byte(utf8);
+            decoded_text += unicode_cpt_to_byte(cpt);
         } catch (const std::out_of_range & /*e*/) {
+            const auto utf8 = unicode_cpt_to_utf8(cpt);
             decoded_text += "[UNK_BYTE_0x";
             for (const auto c : utf8) {
                 decoded_text += format("%02x", (uint8_t) c);

src/models/deepseek4.cpp

@@ -683,6 +683,13 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask(
     GGML_ASSERT(kq_mask);
     GGML_ASSERT(top_k);
 
+    if (cparams.fused_top_k_mask) {
+        ggml_tensor * kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, nullptr, nullptr, nullptr, il);
+        cb(kq_mask_top_k, name, il);
+
+        return kq_mask_top_k;
+    }
+
     ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);
     kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3],
             kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);

src/models/glm5-next.cpp

@@ -323,13 +323,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
         const int64_t n_sel      = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0);
         inp->n_sel = (uint32_t) n_sel;
 
-        // the scatter maps the dead slots of this mask to dump rows
-        inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
-        ggml_set_input(inp->sel_mask);
-        // set_input_kpool always fills the mask, so it stays allocated in every graph
-        ggml_build_forward_expand(gf, inp->sel_mask);
-        // one view for all layers: each view of a host input is a separate upload
-        inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
+        // the scatter maps the dead slots of this mask to dump rows, the fused mask op reads liveness from the scores instead
+        if (!cparams.fused_top_k_mask) {
+            inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
+            ggml_set_input(inp->sel_mask);
+            // set_input_kpool always fills the mask, so it stays allocated in every graph
+            ggml_build_forward_expand(gf, inp->sel_mask);
+            // one view for all layers: each view of a host input is a separate upload
+            inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens);
+        }
     }
 
     inp->n_new = n_new;
@@ -973,6 +975,17 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
         const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
         ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED
 
+        // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere
+        if (cparams.fused_top_k_mask) {
+            ggml_tensor * tail = hparams.indexer_kpool_select_tail ? inp_kpool->tail_idxs : nullptr;
+            GGML_ASSERT(kpool*n_top_pool + (tail ? tail->ne[0] : 0) == (int64_t) inp_kpool->n_sel);
+
+            ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, tail, il);
+            cb(sel, "indexer_sel", il);
+
+            return sel;
+        }
+
         // The selection mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score.
         ggml_tensor * sel_score = ggml_get_rows(ctx0,
                 ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens]

src/models/qwen4exp.cpp

@@ -813,6 +813,16 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel(
     ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], unordered
     cb(top_k, "indexer_top_k", il);
 
+    // keep the kq_mask cells of the live picked pools and the tail, -inf elsewhere
+    if (cparams.fused_top_k_mask) {
+        GGML_ASSERT(kpool*n_top_pool + inp_kpool->tail_idxs->ne[0] == inp_kpool->n_sel);
+
+        ggml_tensor * sel = build_fused_top_k_mask(kq_mask, top_k, score, inp_kpool->pool_idxs, inp_kpool->tail_idxs, il);
+        cb(sel, "indexer_sel", il);
+
+        return sel;
+    }
+
     // the top blocks, then the incomplete tail with n_kv for missing cells
     ggml_tensor * sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs,
             ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens]

src/unicode.cpp

@@ -202,6 +202,19 @@ static std::unordered_map<std::string, uint8_t> unicode_utf8_to_byte_map() {
     return map;
 }
 
+static std::vector<int> unicode_cpt_to_byte_map() {
+    std::vector<int> map;
+    for (int ch = 0; ch < 256; ++ch) {
+        size_t offset = 0;
+        const uint32_t cpt = unicode_cpt_from_utf8(unicode_byte_to_utf8(ch), offset);
+        if (cpt >= map.size()) {
+            map.resize(cpt + 1, -1);
+        }
+        map[cpt] = ch;
+    }
+    return map;
+}
+
 static std::vector<std::string> unicode_byte_encoding_process(const std::vector<std::string> & bpe_words) {
     std::vector<std::string> bpe_encoded_words;
     for (const auto & word : bpe_words) {
@@ -1416,6 +1429,14 @@ uint8_t unicode_utf8_to_byte(const std::string & utf8) {
     return map.at(utf8);
 }
 
+uint8_t unicode_cpt_to_byte(uint32_t cpt) {
+    static const std::vector<int> map = unicode_cpt_to_byte_map();
+    if (cpt >= map.size() || map[cpt] < 0) {
+        throw std::out_of_range("unicode_cpt_to_byte: not a byte token code point");
+    }
+    return map[cpt];
+}
+
 uint32_t unicode_tolower(uint32_t cpt) {
     // binary search
     auto it = std::lower_bound(unicode_map_lowercase.begin(), unicode_map_lowercase.end(), cpt,

src/unicode.h

@@ -104,6 +104,7 @@ unicode_cpt_flags unicode_cpt_flags_from_utf8(const std::string & utf8);
 
 std::string unicode_byte_to_utf8(uint8_t byte);
 uint8_t     unicode_utf8_to_byte(const std::string & utf8);
+uint8_t     unicode_cpt_to_byte (uint32_t cpt);
 
 uint32_t unicode_tolower(uint32_t cpt);
 

tests/CMakeLists.txt

@@ -332,6 +332,7 @@ if (NOT GGML_BACKEND_DL AND GGML_CPU)
     llama_build_and_test(test-rope.cpp LINK ggml)
     llama_build_and_test(test-col2im-1d.cpp LINK ggml)
     llama_build_and_test(test-conv-1d.cpp LINK ggml)
+    llama_build_and_test(test-gelu-quick.cpp LINK ggml)
 endif()
 
 # libmtmd

tests/test-backend-ops.cpp

@@ -2069,6 +2069,15 @@ struct test_unary : public test_case {
         return VARS_TO_STR3(type, ne_a, v);
     }
 
+    double max_nmse_err(ggml_backend_t backend) override {
+        // Metal computes f32 gelu_quick in f32, so an fp16 table in the CPU reference (NMSE about 5e-8) fails this
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (op == GGML_UNARY_OP_GELU_QUICK && type == GGML_TYPE_F32 && strcmp(ggml_backend_reg_name(reg), "MTL") == 0) {
+            return 1e-10;
+        }
+        return test_case::max_nmse_err(backend);
+    }
+
     test_unary(ggml_unary_op op,
             ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {128, 2, 2, 2},
@@ -2157,6 +2166,15 @@ struct test_glu : public test_case {
         return VARS_TO_STR4(type, ne_a, v, swapped);
     }
 
+    double max_nmse_err(ggml_backend_t backend) override {
+        // Metal computes f32 geglu_quick in f32, so an fp16 table in the CPU reference (NMSE about 5e-8) fails this
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (op == GGML_GLU_OP_GEGLU_QUICK && type == GGML_TYPE_F32 && strcmp(ggml_backend_reg_name(reg), "MTL") == 0) {
+            return 1e-10;
+        }
+        return test_case::max_nmse_err(backend);
+    }
+
     test_glu(ggml_glu_op op,
             ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {128, 2, 2, 2},
@@ -2202,6 +2220,15 @@ struct test_glu_split : public test_case {
         return VARS_TO_STR3(type, ne_a, v) + ",split";
     }
 
+    double max_nmse_err(ggml_backend_t backend) override {
+        // Metal computes f32 geglu_quick in f32, so an fp16 table in the CPU reference (NMSE about 5e-8) fails this
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        if (op == GGML_GLU_OP_GEGLU_QUICK && type == GGML_TYPE_F32 && strcmp(ggml_backend_reg_name(reg), "MTL") == 0) {
+            return 1e-10;
+        }
+        return test_case::max_nmse_err(backend);
+    }
+
     test_glu_split(ggml_glu_op op,
             ggml_type type = GGML_TYPE_F32,
             std::array<int64_t, 4> ne_a = {128, 2, 2, 2},
@@ -9101,6 +9128,148 @@ struct test_lightning_indexer : public test_case {
     }
 };
 
+// GGML_OP_TOP_K_MASK
+struct test_top_k_mask : public test_case {
+    const ggml_type type;
+    const int64_t n_kv;
+    const int64_t n_blk;   // 0: no blocks, a pick is a cell
+    const int64_t n_pick;  // number of blocks, ignored without blocks
+    const int64_t n_top;
+    const int64_t n_extra; // 0: no extra cells
+    const bool    score;   // false: every pick is live
+    const int64_t nb;      // rows per stream
+    const int64_t nh;      // rows of top_k per mask row, broadcast over ne2
+    const int64_t ns;      // streams
+    const bool    real;    // blocks of consecutive cells, live picks, as in a real cache
+
+    std::string vars() override {
+        return VARS_TO_STR11(type, n_kv, n_blk, n_pick, n_top, n_extra, score, nb, nh, ns, real);
+    }
+
+    // the result is a copy of the mask or -inf, the harness checks the -inf pattern, the finite values must match exactly
+    double max_err() override {
+        return 0.0;
+    }
+
+    double err(const float * a, const float * b, size_t n) override {
+        double result = 0.0;
+        for (size_t i = 0; i < n; ++i) {
+            if (std::isfinite(a[i]) && std::isfinite(b[i])) {
+                result = std::max(result, (double) fabsf(a[i] - b[i]));
+            }
+        }
+        return result;
+    }
+
+    size_t op_size(ggml_tensor * t) override {
+        const size_t ts = ggml_type_size(type);
+        const size_t nc = std::max<int64_t>(n_blk, 1);
+        return ggml_nbytes(t) + nb*nh*ns*(n_top*(sizeof(int32_t) + (score ? sizeof(float) : 0) + nc*(sizeof(int32_t) + 2*ts)) + n_extra*(sizeof(int32_t) + 2*ts));
+    }
+
+    test_top_k_mask(ggml_type type = GGML_TYPE_F16, int64_t n_kv = 256, int64_t n_blk = 4, int64_t n_pick = 64,
+            int64_t n_top = 16, int64_t n_extra = 3, bool score = true, int64_t nb = 32, int64_t nh = 1, int64_t ns = 1, bool real = false)
+        : type(type), n_kv(n_kv), n_blk(n_blk), n_pick(n_blk ? n_pick : n_kv), n_top(n_top), n_extra(n_extra), score(score),
+          nb(nb), nh(nh), ns(ns), real(real) {}
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * m = ggml_new_tensor_4d(ctx, type, n_kv, nb, 1, ns);
+        ggml_set_name(m, "m");
+
+        ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top, nb, nh, ns);
+        ggml_set_name(top_k, "top_k");
+
+        ggml_tensor * s = nullptr;
+        if (score) {
+            s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_pick, nb, nh, ns);
+            ggml_set_name(s, "score");
+        }
+
+        ggml_tensor * blocks = nullptr;
+        if (n_blk > 0) {
+            blocks = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_blk, n_pick);
+            ggml_set_name(blocks, "blocks");
+        }
+
+        ggml_tensor * extra = nullptr;
+        if (n_extra > 0) {
+            extra = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_extra, nb, nh, ns);
+            ggml_set_name(extra, "extra");
+        }
+
+        ggml_tensor * out = ggml_top_k_mask(ctx, m, top_k, s, blocks, extra);
+        ggml_set_name(out, "out");
+
+        return out;
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        std::mt19937 gen(0x4B50);
+
+        // cells and picks sometimes fall out of range, the op skips them
+        auto init_cells = [&](ggml_tensor * t) {
+            std::vector<int32_t> data(ggml_nelements(t));
+            if (real) {
+                // block p holds cells p*n_blk..p*n_blk + n_blk - 1, the extra cells are the last cells
+                const bool is_blocks = strcmp(t->name, "blocks") == 0;
+                for (int64_t i = 0; i < (int64_t) data.size(); ++i) {
+                    const int64_t c = is_blocks ? i : n_kv - 1 - i % t->ne[0];
+                    data[i] = c < n_kv ? (int32_t) c : (int32_t) n_kv;
+                }
+            } else {
+                std::uniform_int_distribution<int32_t> cell(0, n_kv - 1);
+                for (auto & c : data) {
+                    const uint32_t r = gen() % 16;
+                    c = r == 0 ? (int32_t) n_kv : r == 1 ? -1 : cell(gen);
+                }
+            }
+            ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t));
+        };
+
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (strcmp(t->name, "m") == 0) {
+                if (type == GGML_TYPE_F16) {
+                    init_tensor_kq_mask(t);
+                } else {
+                    std::vector<float> data(ggml_nelements(t));
+                    std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+                    for (auto & v : data) {
+                        v = gen() % 4 == 0 ? -INFINITY : dis(gen);
+                    }
+                    ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float));
+                }
+            } else if (strcmp(t->name, "score") == 0) {
+                // -inf marks the picks a row cannot see
+                std::vector<float> data(ggml_nelements(t));
+                std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+                for (auto & v : data) {
+                    v = !real && gen() % 4 == 0 ? -INFINITY : dis(gen);
+                }
+                ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float));
+            } else if (strcmp(t->name, "top_k") == 0) {
+                std::vector<int32_t> data(ggml_nelements(t));
+                std::vector<int32_t> order(n_pick);
+                for (int64_t r = 0; r < ggml_nrows(t); ++r) {
+                    for (int64_t i = 0; i < n_pick; ++i) {
+                        order[i] = i;
+                    }
+                    std::shuffle(order.begin(), order.end(), gen);
+                    for (int64_t j = 0; j < n_top; ++j) {
+                        const bool out = !real && j == 0 && r % 5 == 0;
+                        data[r*n_top + j] = out ? (r % 2 ? -1 : (int32_t) n_pick) : order[j % n_pick];
+                    }
+                }
+                ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t));
+            } else if (strcmp(t->name, "blocks") == 0 || strcmp(t->name, "extra") == 0) {
+                init_cells(t);
+            } else {
+                // the harness sentinels and the output, the callback checks the sentinels for NaN
+                init_tensor_uniform(t);
+            }
+        }
+    }
+};
+
 // Deserializable generic test case
 struct input_tensor {
     ggml_type type;
@@ -9885,6 +10054,25 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_dsv4_hc_post(2560, 21, true, true));
     test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, true));
 
+    for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
+        // blocks with extra cells and scores (k-pool)
+        test_cases.emplace_back(new test_top_k_mask(type));
+        test_cases.emplace_back(new test_top_k_mask(type, 1001, 4, 70, 16, 3, true, 7));
+        test_cases.emplace_back(new test_top_k_mask(type, 256, 4, 64, 64, 0, true, 33));
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 8, 16, 32, 7, true, 5, 1, 4));
+        test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 4096, 4, 1024, 512, 3, true, 64, 1, 2));
+        // cells without scores (DSA)
+        test_cases.emplace_back(new test_top_k_mask(type, 256, 0, 0, 16, 0, false, 32));
+        test_cases.emplace_back(new test_top_k_mask(type, 1001, 0, 0, 64, 0, false, 7, 1, 3));
+        test_cases.emplace_back(new test_top_k_mask(type, 4096, 0, 0, 2048, 0, false, 5));
+        // cells with scores and extra cells
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 0, 0, 32, 5, true, 9));
+        // blocks picked per head, the mask broadcast over the heads
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 16, 32, 4, 0, true, 6, 4, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 1024, 8, 128, 16, 8, false, 3, 2, 2));
+    }
+
     // glu ops
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) {
         for (int v : {0, 1}) {
@@ -12718,6 +12906,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // top_k_mask, Flash-Next QSA shapes: 512 pools of 4 + a 3-cell tail, and DSA shapes: 2048 cells
+    for (int64_t n_kv : { 18432, 131072 }) {
+        for (int64_t nb : { 1, 2048 }) {
+            test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 4, n_kv/4, 512, 3, true, nb, 1, 1, true));
+            test_cases.emplace_back(new test_top_k_mask(GGML_TYPE_F16, n_kv, 0, 0, 2048, 0, false, nb, 1, 1, true));
+        }
+    }
+
     // launch-overhead isolation: single L2_NORM launch vs batched siblings at the GDN
     // production shape (strided qkv views) -- perf-mode only, the eval list has its own
     // 2/4-wide coverage

tests/test-gelu-quick.cpp

@@ -0,0 +1,131 @@
+// test-gelu-quick.cpp: validate the F32 gelu_quick and geglu_quick CPU kernels against the exact formula.
+//
+// test-backend-ops uses the CPU backend as its reference, so it cannot see an error in the CPU kernels
+// themselves. This test computes gelu_quick(x) = x / (1 + exp(-1.702 x)) in double precision and compares
+// every F32 output of the CPU backend with it: dense ranges, values outside the F16 range, values near
+// the float limits, and lengths that leave a tail after the vector loop.
+
+#include "ggml.h"
+#include "ggml-cpu.h"
+
+#include <cfloat>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+// Same coefficient as the CPU kernels
+static const double GELU_QUICK_COEF_REF = (double) 1.702f;
+
+static double gelu_quick_ref(double x) {
+    return x / (1.0 + std::exp(-GELU_QUICK_COEF_REF * x));
+}
+
+// Deterministic LCG mapped to [-1, 1]
+static uint64_t g_rng = 0x12345678ULL;
+static float frand(void) {
+    g_rng = g_rng * 6364136223846793005ULL + 1442695040888963407ULL;
+    return (float)((g_rng >> 33) & 0xffffff) / (float)0x800000 - 1.0f;
+}
+
+static std::vector<float> make_inputs(void) {
+    std::vector<float> x;
+    for (int i = 0; i <= 200000; i++) {
+        x.push_back(-10.0f + 20.0f * (float) i / 200000.0f);
+    }
+    for (int i = 0; i <= 2000; i++) {
+        x.push_back(-150.0f + 300.0f * (float) i / 2000.0f);
+    }
+    // outside the F16 range, near the float limits, and subnormal; both signs
+    const float edge[] = { 65504.0f, 65520.0f, 1e5f, 1e10f, 1e30f, FLT_MAX, 1e-30f, FLT_MIN, FLT_MIN / 16, 0.0f };
+    for (float v : edge) {
+        x.push_back(v);
+        x.push_back(-v);
+    }
+    return x;
+}
+
+static bool matches(float y, double ref) {
+    if (std::isnan(ref)) {
+        return std::isnan(y);
+    }
+    if (std::isinf(ref)) {
+        return (double) y == ref;
+    }
+    return std::fabs((double) y - ref) <= 1e-6 + 1e-6 * std::fabs(ref);
+}
+
+// Runs gelu_quick and geglu_quick on x[0, n) and returns the number of outputs that differ from the reference
+static int check(const float * x, const float * g, int64_t n) {
+    struct ggml_init_params params = {
+        /* .mem_size   = */ (size_t) 64 << 20,
+        /* .mem_base   = */ NULL,
+        /* .no_alloc   = */ false,
+    };
+    struct ggml_context * ctx = ggml_init(params);
+
+    struct ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n);
+    struct ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n);
+    memcpy(a->data, x, n * sizeof(float));
+    memcpy(b->data, g, n * sizeof(float));
+
+    struct ggml_tensor * y_gelu  = ggml_gelu_quick(ctx, a);
+    struct ggml_tensor * y_geglu = ggml_geglu_quick_split(ctx, a, b);
+
+    struct ggml_cgraph * gf = ggml_new_graph(ctx);
+    ggml_build_forward_expand(gf, y_gelu);
+    ggml_build_forward_expand(gf, y_geglu);
+    ggml_graph_compute_with_ctx(ctx, gf, 1);
+
+    int    bad     = 0;
+    double max_err = 0.0;
+    for (int64_t i = 0; i < n; i++) {
+        const double r0 = gelu_quick_ref(x[i]);
+        const double r1 = r0 * g[i];
+        const float  y0 = ((const float *) y_gelu->data)[i];
+        const float  y1 = ((const float *) y_geglu->data)[i];
+        if (std::isfinite(r0) && std::isfinite(y0)) {
+            max_err = std::fmax(max_err, std::fabs((double) y0 - r0) / (1.0 + std::fabs(r0)));
+        }
+        if (!matches(y0, r0)) {
+            if (bad < 3) {
+                printf("  gelu_quick(%g) = %g, expected %g\n", x[i], y0, r0);
+            }
+            bad++;
+        }
+        if (!matches(y1, r1)) {
+            if (bad < 3) {
+                printf("  geglu_quick(%g, %g) = %g, expected %g\n", x[i], g[i], y1, r1);
+            }
+            bad++;
+        }
+    }
+    printf("gelu_quick n=%6d: max err %.2e, %d mismatches %s\n", (int) n, max_err, bad, bad == 0 ? "OK" : "FAIL");
+
+    ggml_free(ctx);
+    return bad;
+}
+
+int main(void) {
+    const std::vector<float> x = make_inputs();
+    std::vector<float> g(x.size());
+    for (float & v : g) {
+        v = frand();
+    }
+
+    int fails = 0;
+
+    // every input in one call
+    fails += check(x.data(), g.data(), (int64_t) x.size()) != 0;
+
+    // short lengths end on the edge values and leave a tail after the vector loop
+    const int64_t lengths[] = { 1, 2, 3, 5, 7, 8, 9, 15, 16, 17, 23, 31, 33 };
+    for (int64_t n : lengths) {
+        const int64_t off = (int64_t) x.size() - n;
+        fails += check(x.data() + off, g.data() + off, n) != 0;
+    }
+
+    printf(fails == 0 ? "all gelu_quick checks passed\n" : "%d gelu_quick checks FAILED\n", fails);
+    return fails == 0 ? 0 : 1;
+}

tests/test-jinja.cpp

@@ -25,6 +25,7 @@ static void test_expressions(testing & t);
 static void test_set_statement(testing & t);
 static void test_filters(testing & t);
 static void test_literals(testing & t);
+static void test_hexadecimal_escapes(testing & t);
 static void test_comments(testing & t);
 static void test_macros(testing & t);
 static void test_namespace(testing & t);
@@ -65,6 +66,7 @@ int main(int argc, char *argv[]) {
     t.test("set statement", test_set_statement);
     t.test("filters", test_filters);
     t.test("literals", test_literals);
+    t.test("hexadecimal escapes", test_hexadecimal_escapes);
     t.test("comments", test_comments);
     t.test("macros", test_macros);
     t.test("namespace", test_namespace);
@@ -1171,6 +1173,54 @@ static void test_literals(testing & t) {
     );
 }
 
+static void test_hexadecimal_escapes(testing & t) {
+    test_template(t, "remove NUL from message content",
+        R"({{ message.content | replace("\x00", "") }})",
+        {{"message", {{"content", std::string("\0a\0b\0", 5)}}}},
+        "ab"
+    );
+
+    test_template(t, "non-ASCII hex escape produces UTF-8",
+        R"({{ '\xE9' }})",
+        json::object(),
+        "\xc3\xa9"
+    );
+
+    test_template(t, "consume exactly two hex digits",
+        R"({{ '\x414' }})",
+        json::object(),
+        "A4"
+    );
+
+    test_template(t, "escaped backslashes remain literal",
+        R"({{ '\\xE9' }})",
+        json::object(),
+        R"(\xE9)"
+    );
+
+    if (!g_python_mode) {
+        const char * invalid_templates[] = {
+            R"({{ '\x' }})",
+            R"({{ '\x0' }})",
+            R"({{ '\xG1' }})",
+            R"({{ '\x1G' }})",
+            R"({{ '\x)",
+            R"({{ '\x0)",
+        };
+        for (const auto * tmpl : invalid_templates) {
+            t.test("reject " + json(tmpl).dump(), [&](testing & t) {
+                bool rejected = false;
+                try {
+                    jinja::lexer().tokenize(tmpl);
+                } catch (const jinja::lexer_exception &) {
+                    rejected = true;
+                }
+                t.assert_true("malformed hex escape must raise a lexer error", rejected);
+            });
+        }
+    }
+}
+
 static void test_comments(testing & t) {
     test_template(t, "inline comment",
         "before{# comment #}after",

tools/mtmd/README-dev.md

@@ -55,7 +55,7 @@ For example, Qwen3-TTS:
 Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system is designed to be flexible and reusable by new models.
 
 `mtmd_gen_audio` is split into 2 main API:
-- Core API `mtmd.h`: handles main inference. Important: the API surface must be stateless; caller must handle state management and audio frame accumulation.
+- Core API `mtmd.h`: handles main inference. Important: `mtmd_gen_state` holds only the per-sequence data that the mmproj itself writes and reads back (RNG, decoder KV / conv state, output buffers). Data that goes to or comes from the backbone (hidden states, codes, PCM) must stay explicit in `mtmd_gen_inp` / `mtmd_gen_out`; the caller drives the backbone and handles audio frame accumulation.
 - Helper API `mtmd-helper.h`: provides a model-agnostic stateful API. Usage example can be found in the `tools/tts` directory.
 
 ### Checklist for porting new audio generation models to mtmd

tools/mtmd/clip-impl.h

@@ -16,6 +16,7 @@
 #include <vector>
 #include <memory>
 #include <fstream>
+#include <random>
 
 #ifdef _WIN32
 #ifndef NOMINMAX
@@ -463,6 +464,15 @@ struct clip_enc_state {
     std::vector<ggml_tensor *> v;
 };
 
+// per-sequence state of audio generation, see clip_gen_state_init()
+struct clip_gen_state {
+    std::mt19937 rng;
+
+    ggml_context_ptr        ctx_data;
+    ggml_backend_buffer_ptr buf;
+    std::vector<ggml_tensor *> slots; // same order as list_gen_state_slots(), F32 [ne0, ne1]
+};
+
 // forward declaration
 // TODO: improve this later
 struct clip_ctx;

tools/mtmd/clip.cpp

@@ -174,10 +174,6 @@ struct clip_ctx {
 
     bool support_batch = false;
 
-    // for audio gen, reseeded only when the caller asks for another seed
-    std::mt19937 rng{std::random_device{}()};
-    uint32_t rng_seed = UINT32_MAX;
-
     clip_ctx(clip_context_params & ctx_params) {
         flash_attn_type = ctx_params.flash_attn_type;
         no_alloc = ctx_params.no_alloc;
@@ -4717,6 +4713,59 @@ static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hpa
     }
 }
 
+// state_out can have another shape than its slot, only the byte size must match
+static void copy_gen_state_out(ggml_tensor * src, ggml_tensor * dst) {
+    GGML_ASSERT(ggml_nbytes(src) == ggml_nbytes(dst));
+    if (src->type == dst->type && ggml_are_same_shape(src, dst) && ggml_is_contiguous(src)) {
+        ggml_backend_tensor_copy(src, dst);
+        return;
+    }
+    std::vector<uint8_t> tmp(ggml_nbytes(src));
+    ggml_backend_tensor_get(src, tmp.data(), 0, tmp.size());
+    ggml_backend_tensor_set(dst, tmp.data(), 0, tmp.size());
+}
+
+clip_gen_state * clip_gen_state_init(clip_ctx * ctx, uint32_t seed) {
+    auto st = std::make_unique<clip_gen_state>();
+    st->rng.seed(seed == UINT32_MAX ? std::random_device{}() : seed);
+
+    const auto slots = list_gen_state_slots(ctx->model.hparams, ctx->model);
+    if (slots.empty()) {
+        return st.release();
+    }
+
+    ggml_init_params params = {
+        /*.mem_size   =*/ ggml_tensor_overhead() * slots.size(),
+        /*.mem_buffer =*/ nullptr,
+        /*.no_alloc   =*/ true,
+    };
+    st->ctx_data.reset(ggml_init(params));
+    for (const auto & slot : slots) {
+        ggml_tensor * t = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F32, slot.ne0, slot.ne1);
+        ggml_format_name(t, "gen_state_%s", slot.name.c_str());
+        st->slots.push_back(t);
+    }
+    st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), ctx->backend));
+    if (!st->buf) {
+        LOG_ERR("%s: failed to allocate the gen state\n", __func__);
+        return nullptr;
+    }
+    // a zero state is the cold start
+    ggml_backend_buffer_clear(st->buf.get(), 0);
+
+    return st.release();
+}
+
+void clip_gen_state_free(clip_gen_state * state) {
+    delete state;
+}
+
+void clip_gen_state_reset(clip_gen_state * state) {
+    if (state->buf) {
+        ggml_backend_buffer_clear(state->buf.get(), 0);
+    }
+}
+
 clip_enc_state * clip_enc_state_init(clip_ctx * ctx, int32_t n_frames_max) {
     const auto & hparams = ctx->model.hparams;
     if (ctx->proj_type() != PROJECTOR_TYPE_VOXTRAL_RT || hparams.attn_window_size <= 0) {
@@ -4782,9 +4831,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
         clip_model_loader::warmup(*ctx, *params->imgs);
     }
 
-    if (params->seed != ctx->rng_seed) {
-        ctx->rng_seed = params->seed;
-        ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed);
+    if (params->gen_process != CLIP_GEN_PROCESS_GEN_UNKNOWN && params->gen_state == nullptr) {
+        LOG_ERR("%s: audio gen requires a gen state\n", __func__);
+        return false;
     }
 
     // build the inference graph
@@ -4848,19 +4897,13 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
         ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur));
     };
 
-    // upload the decoder state from the previous call, or zero-fill on a cold start
+    // state of the previous call, zero on a cold start
     auto set_gen_state_in = [&]() {
-        size_t offset = 0;
+        const auto & slots = params->gen_state->slots;
+        size_t i = 0;
         for (const auto & slot : list_gen_state_slots(hparams, model)) {
-            ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str());
-            const size_t nb = ggml_nbytes(t);
-            if (params->state_in && params->state_in->size() >= offset + nb) {
-                ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb);
-            } else {
-                std::vector<uint8_t> zeros(nb, 0);
-                ggml_backend_tensor_set(t, zeros.data(), 0, nb);
-            }
-            offset += nb;
+            GGML_ASSERT(i < slots.size());
+            ggml_backend_tensor_copy(slots[i++], get_inp_tensor(("state_in_" + slot.name).c_str()));
         }
     };
 
@@ -5584,7 +5627,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                     std::normal_distribution<float> dist(0.0f, std::sqrt(temp));
                     std::vector<float> noise(ggml_nelements(t));
                     for (auto & v : noise) {
-                        v = dist(ctx->rng);
+                        v = dist(params->gen_state->rng);
                     }
                     set_input_f32("inp_noise", noise);
                 }
@@ -5732,7 +5775,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
                     std::uniform_real_distribution<float> dist(0.0f, 1.0f);
                     const int64_t n_acoustic = model.gen_code_head_w->ne[2];
                     for (int64_t g = 0; g < n_acoustic; g++) {
-                        std::vector<float> r = { dist(ctx->rng) };
+                        std::vector<float> r = { dist(params->gen_state->rng) };
                         set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r);
                     }
                 }
@@ -6262,22 +6305,16 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
         st->n_past += n_new;
     }
 
-    if (params->state_out != nullptr) {
-        auto & state_out = *params->state_out;
-        size_t total = 0;
-        for (const auto & slot : list_gen_state_slots(hparams, model)) {
-            total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float);
-        }
-        state_out.resize(total);
-        size_t offset = 0;
+    if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
+        const auto & slots = params->gen_state->slots;
+        size_t i = 0;
         for (const auto & slot : list_gen_state_slots(hparams, model)) {
             ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str());
             if (t == nullptr) {
-                GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str());
+                GGML_ABORT("graph has no \"state_out_%s\" tensor", slot.name.c_str());
             }
-            const size_t nb = ggml_nbytes(t);
-            ggml_backend_tensor_get(t, state_out.data() + offset, 0, nb);
-            offset += nb;
+            GGML_ASSERT(i < slots.size());
+            copy_gen_state_out(t, slots[i++]);
         }
     }
 

tools/mtmd/clip.h

@@ -14,6 +14,7 @@
 
 struct clip_ctx;
 struct clip_enc_state;
+struct clip_gen_state;
 
 struct clip_image_size {
     int width;
@@ -110,7 +111,6 @@ struct clip_encode_params {
     float   top_p = 1.0f;
     std::vector<int32_t> * out_codes = nullptr; // this frame's 16 sampled codes
     std::vector<float> * out_feats = nullptr; // continuous counterpart of out_codes
-    uint32_t seed = UINT32_MAX;               // UINT32_MAX for random
     float   temp = 0.0f;                      // sampling temperature, noise scale for flow-matching decoders
     bool * out_is_eos = nullptr;
 
@@ -118,8 +118,9 @@ struct clip_encode_params {
     const std::vector<int32_t> * codes = nullptr;     // this frame's 16 RVQ codes
     const std::vector<float> *   feats = nullptr;     // continuous counterpart of codes
     std::vector<float> * out_audio = nullptr;         // decoded PCM samples, F32
-    const std::vector<uint8_t> * state_in  = nullptr; // state from previous call, null or wrong size means cold start
-    std::vector<uint8_t> *       state_out = nullptr; // state for the next call
+
+    // audio gen: state of one sequence, clip_encode() updates it in place
+    struct clip_gen_state * gen_state = nullptr;
 
     // streaming audio input, updated in place by clip_encode()
     struct clip_enc_state * enc_state = nullptr;
@@ -132,6 +133,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params);
 struct clip_enc_state * clip_enc_state_init(struct clip_ctx * ctx, int32_t n_frames_max);
 void                    clip_enc_state_free(struct clip_enc_state * state);
 
+// per-sequence state of audio generation: RNG and the slots kept between GEN_WAV calls
+// seed UINT32_MAX means random; reset() clears the slots and keeps the RNG stream
+struct clip_gen_state * clip_gen_state_init(struct clip_ctx * ctx, uint32_t seed);
+void                    clip_gen_state_free(struct clip_gen_state * state);
+void                    clip_gen_state_reset(struct clip_gen_state * state);
+
 bool clip_is_llava(const struct clip_ctx * ctx);
 // note for contributor: this clip_is_(model) pattern is deprecated
 //                       do NOT add new functions like this

tools/mtmd/mtmd-helper-gen.cpp

@@ -105,12 +105,25 @@ public:
     virtual int32_t flush() = 0;
 
 protected:
+    // one state per generation, so a fixed seed always gives the same output
+    bool init_gen_state(uint32_t seed) {
+        mtmd_gen_state_params params = mtmd_gen_state_params_default();
+        params.seed = seed;
+        gen_state.reset(mtmd_gen_state_init(mctx, params));
+        if (!gen_state) {
+            LOG_ERR("mtmd_helper_gen_audio: failed to create the gen state\n");
+            return false;
+        }
+        return true;
+    }
+
     llama_context * lctx;
     mtmd_context  * mctx;
     const llama_model * model;
     const llama_vocab  * vocab;
     int n_embd;
     mtmd_gen_audio_info info;
+    mtmd::gen_state_ptr gen_state;
 };
 
 // Qwen3-TTS: backbone samples codec_0, code_predictor gives the other 15 codebooks,
@@ -123,7 +136,7 @@ public:
         seq_id = 0;
         pos = 0;
         codes_buf.clear();
-        c2w_state.clear();
+        gen_state.reset();
         audio_pcm.clear();
         overlay.clear();
         h_state_buf.clear();
@@ -220,9 +233,11 @@ public:
         const mtmd_gen_inp def = mtmd_gen_inp_default(mctx);
         top_k = inp->top_k > 0 ? inp->top_k : def.top_k;
         top_p = inp->top_p > 0 ? inp->top_p : def.top_p;
-        seed  = inp->seed;
         out_type = inp->out_type;
         stream = inp->stream;
+        if (!init_gen_state(inp->seed)) {
+            return 1;
+        }
 
         // the prompt above holds the whole text stream up to tts_eos, so every generated
         // frame adds tts_pad on top of the codes embedding
@@ -279,9 +294,8 @@ public:
         inp.embd  = const_cast<float *>(h_state_in);
         inp.top_k = top_k;
         inp.top_p = top_p;
-        inp.seed  = seed;
         mtmd_gen_out out{};
-        if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
+        if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
             LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n");
             return 1;
         }
@@ -449,16 +463,12 @@ private:
         inp.type       = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
         inp.codes      = codes_buf.data();
         inp.n_codes    = codes_buf.size();
-        inp.seed       = seed; // same seed as gen_code, else clip reseeds mid-generation
-        inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data();
-        inp.state_size = c2w_state.size();
         mtmd_gen_out out{};
-        if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
+        if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
             LOG_ERR("mtmd_helper_gen_audio: gen_wav process failed\n");
             return false;
         }
         audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
-        c2w_state.assign(out.state_data, out.state_data + out.state_size);
         codes_buf.clear();
         return true;
     }
@@ -491,9 +501,7 @@ private:
     int prompt_pos = 0;
     int32_t  top_k = 50;
     float    top_p = 1.0f;
-    uint32_t seed  = UINT32_MAX;
     std::vector<int32_t> codes_buf;
-    std::vector<uint8_t> c2w_state;
     std::vector<float>   audio_pcm;
     std::vector<float> overlay;
     std::vector<float> h_state_buf;
@@ -539,7 +547,7 @@ public:
         seq_id = 0;
         pos = 0;
         feats_buf.clear();
-        dec_state.clear();
+        gen_state.reset();
         audio_pcm.clear();
         h_state_buf.clear();
         out_buf.clear();
@@ -623,9 +631,11 @@ public:
         prompt_batch->set_position_normal(0, seq_id);
         prompt_pos = 0;
 
-        seed     = inp->seed;
         out_type = inp->out_type;
         stream   = inp->stream;
+        if (!init_gen_state(inp->seed)) {
+            return 1;
+        }
 
         return 0;
     }
@@ -663,13 +673,11 @@ public:
         mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
         inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
         inp.embd = const_cast<float *>(h_state_in);
-        // clip only reseeds when the seed changes, so pass the same one on every step
-        inp.seed = seed;
         if (pack.temp > 0.0f) {
             inp.temp = pack.temp;
         }
         mtmd_gen_out out{};
-        if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
+        if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
             LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n");
             return 1;
         }
@@ -895,7 +903,7 @@ private:
             return 1;
         }
         // the decoder restarts too, the next chunk's audio is not continuous with this one
-        dec_state.clear();
+        mtmd_gen_state_reset(gen_state.get());
 
         if (chunk_idx + 1 >= chunks.size()) {
             *out_stop    = true;
@@ -1023,16 +1031,12 @@ private:
         inp.type       = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
         inp.feats      = feats_buf.data();
         inp.n_feats    = feats_buf.size();
-        inp.seed       = seed;
-        inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data();
-        inp.state_size = dec_state.size();
         mtmd_gen_out out{};
-        if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) {
+        if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
             LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n");
             return false;
         }
         audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
-        dec_state.assign(out.state_data, out.state_data + out.state_size);
         feats_buf.clear();
         return true;
     }
@@ -1049,7 +1053,6 @@ private:
     std::unique_ptr<decode_embd_batch> prompt_batch;
     int n_prompt = 0;
     int prompt_pos = 0;
-    uint32_t seed = UINT32_MAX;
     // end-of-speech is latched, then a few more frames are generated as tail padding
     int step_idx = 0;
     int eos_step = -1;
@@ -1064,7 +1067,6 @@ private:
     // latents are decoded a window at a time, the decoder state bridges the windows
     size_t window_frames = 8;
     std::vector<float>   feats_buf;
-    std::vector<uint8_t> dec_state;
     std::vector<float>   audio_pcm;
     std::vector<float>   h_state_buf;
     mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;

tools/mtmd/mtmd.cpp

@@ -523,11 +523,6 @@ struct mtmd_context {
 
     // generation context
     struct clip_ctx * ctx_gen_a; // audio
-    std::vector<int32_t> gen_out_codes; // this frame's 16 sampled codes (GEN_CODE)
-    std::vector<float>   gen_out_feats; // this frame's continuous features, if any (GEN_CODE)
-    std::vector<float>   gen_out_embd;  // next-step hidden state fed back to backbone (GEN_CODE)
-    std::vector<float>   gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV)
-    std::vector<uint8_t> gen_out_state; // state to feed into the next GEN_WAV call
 
     bool print_timings;
     int n_threads;
@@ -2080,7 +2075,6 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
 mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
     mtmd_gen_inp inp{};
     inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
-    inp.seed = UINT32_MAX;
     if (!ctx->ctx_gen_a) {
         return inp;
     }
@@ -2104,12 +2098,58 @@ mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
     return inp;
 }
 
-static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
+struct mtmd_gen_state {
+    clip_ctx * ctx_clip = nullptr;
+    clip_gen_state * clip_state = nullptr;
+
+    std::vector<int32_t> out_codes; // this frame's 16 sampled codes (GEN_CODE)
+    std::vector<float>   out_feats; // this frame's continuous features, if any (GEN_CODE)
+    std::vector<float>   out_embd;  // next-step hidden state fed back to backbone (GEN_CODE)
+    std::vector<float>   out_audio; // decoded PCM samples for the current frame (GEN_WAV)
+
+    ~mtmd_gen_state() {
+        clip_gen_state_free(clip_state);
+    }
+};
+
+mtmd_gen_state_params mtmd_gen_state_params_default(void) {
+    mtmd_gen_state_params params{};
+    params.seed = UINT32_MAX;
+    return params;
+}
+
+mtmd_gen_state * mtmd_gen_state_init(mtmd_context * ctx, mtmd_gen_state_params params) {
+    if (!ctx->ctx_gen_a) {
+        LOG_ERR("%s: model does not support audio generation\n", __func__);
+        return nullptr;
+    }
+    auto st = std::make_unique<mtmd_gen_state>();
+    st->ctx_clip   = ctx->ctx_gen_a;
+    st->clip_state = clip_gen_state_init(ctx->ctx_gen_a, params.seed);
+    if (!st->clip_state) {
+        return nullptr;
+    }
+    return st.release();
+}
+
+void mtmd_gen_state_free(mtmd_gen_state * state) {
+    delete state;
+}
+
+void mtmd_gen_state_reset(mtmd_gen_state * state) {
+    clip_gen_state_reset(state->clip_state);
+}
+
+static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * state, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
     clip_ctx * ctx_clip = ctx->ctx_gen_a;
     if (!ctx_clip) {
         LOG_ERR("%s: model does not support audio generation\n", __func__);
         return 1;
     }
+    if (!state || state->ctx_clip != ctx_clip) {
+        LOG_ERR("%s: invalid gen state\n", __func__);
+        return 1;
+    }
 
     *out = {};
 
@@ -2139,24 +2179,24 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
         params.code0         = inp->code0;
         params.top_k         = inp->top_k;
         params.top_p         = inp->top_p;
-        params.seed          = inp->seed;
         params.temp          = inp->temp;
         params.out_is_eos    = &is_eos;
+        params.gen_state     = state->clip_state;
 
         if (!clip_encode(ctx_clip, &params)) {
             LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__);
             return 1;
         }
 
-        ctx->gen_out_embd  = std::move(out_embd);
-        ctx->gen_out_codes = std::move(out_codes);
-        ctx->gen_out_feats = std::move(out_feats);
+        state->out_embd  = std::move(out_embd);
+        state->out_codes = std::move(out_codes);
+        state->out_feats = std::move(out_feats);
 
-        out->embd      = ctx->gen_out_embd.data();
-        out->codes     = ctx->gen_out_codes.data();
-        out->n_codes   = ctx->gen_out_codes.size();
-        out->feats     = ctx->gen_out_feats.data();
-        out->n_feats   = ctx->gen_out_feats.size();
+        out->embd      = state->out_embd.data();
+        out->codes     = state->out_codes.data();
+        out->n_codes   = state->out_codes.size();
+        out->feats     = state->out_feats.data();
+        out->n_feats   = state->out_feats.size();
         out->is_eos    = is_eos;
         return 0;
     }
@@ -2175,10 +2215,6 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
     } else {
         in_feats.assign(inp->feats, inp->feats + inp->n_feats);
     }
-    std::vector<uint8_t> in_state;
-    if (inp->state_data) {
-        in_state.assign(inp->state_data, inp->state_data + inp->state_size);
-    }
 
     // gen_wav has no hidden-state input, the batch entry is an unused placeholder
     // TODO @ngxson : some models in the future may require hidden-state input, need to update this code later
@@ -2194,30 +2230,25 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in
     params.imgs        = &batch;
     params.n_threads   = ctx->n_threads;
     params.gen_process = CLIP_GEN_PROCESS_GEN_WAV;
-    // gen_wav draws no randomness, but keep the seed so it does not reseed mid-generation
-    params.seed        = inp->seed;
     params.codes       = has_codes ? &in_codes : nullptr;
     params.feats       = has_feats ? &in_feats : nullptr;
-    params.out_audio   = &ctx->gen_out_audio;
-    params.state_in    = inp->state_data ? &in_state : nullptr;
-    params.state_out   = &ctx->gen_out_state;
+    params.out_audio   = &state->out_audio;
+    params.gen_state   = state->clip_state;
 
     if (!clip_encode(ctx_clip, &params)) {
         LOG_ERR("%s: clip_encode failed (code2wav)\n", __func__);
         return 1;
     }
 
-    out->audio      = ctx->gen_out_audio.data();
-    out->n_samples  = ctx->gen_out_audio.size();
-    out->state_data = (const char *) ctx->gen_out_state.data();
-    out->state_size = ctx->gen_out_state.size();
+    out->audio      = state->out_audio.data();
+    out->n_samples  = state->out_audio.size();
 
     return 0;
 }
 
-int32_t mtmd_gen_audio_process(mtmd_context * ctx, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
+int32_t mtmd_gen_audio_process(mtmd_context * ctx, mtmd_gen_state * state, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
     try {
-        return mtmd_gen_audio_process_impl(ctx, inp, out);
+        return mtmd_gen_audio_process_impl(ctx, state, inp, out);
     } catch (const std::exception & e) {
         LOG_ERR("%s: error: %s\n", __func__, e.what());
         return 1;

tools/mtmd/mtmd.h

@@ -65,6 +65,7 @@ struct mtmd_image_tokens;
 struct mtmd_input_chunk;
 struct mtmd_input_chunks;
 struct mtmd_batch;
+struct mtmd_gen_state;
 
 struct mtmd_input_text {
     const char * text;
@@ -79,6 +80,7 @@ struct mtmd_input_part {
     const struct mtmd_bitmap * bitmap;
 };
 
+
 //
 // C API
 //
@@ -91,6 +93,7 @@ typedef struct mtmd_input_chunks mtmd_input_chunks;
 typedef struct mtmd_input_text   mtmd_input_text;
 typedef struct mtmd_input_part   mtmd_input_part;
 typedef struct mtmd_batch        mtmd_batch;
+typedef struct mtmd_gen_state    mtmd_gen_state;
 
 typedef bool (*mtmd_progress_callback)(float progress, void * user_data);
 
@@ -443,8 +446,7 @@ struct mtmd_gen_inp {
     float * embd;   // the hidden state from backbone, must have n_text_embd elements
     int32_t top_k;
     float   top_p;
-    uint32_t seed; // UINT32_MAX for random
-    float    temp; // sampling temperature, or noise scale for flow-matching decoders
+    float   temp; // sampling temperature, or noise scale for flow-matching decoders
 
     // for MTMD_GEN_PROCESS_TYPE_GEN_WAV
     // pass either codes (discrete) or feats (continuous), depending on the pipeline
@@ -452,12 +454,10 @@ struct mtmd_gen_inp {
     size_t    n_codes;
     const float * feats;
     size_t        n_feats;
-    const char * state_data;
-    size_t       state_size;
 };
 
 struct mtmd_gen_out {
-    // note: output memory is allocated by the context, valid until next process() call
+    // note: output memory is owned by the state, valid until the next process() call with this state
 
     // for MTMD_GEN_PROCESS_TYPE_GEN_CODE
     const int32_t * codes;
@@ -471,15 +471,28 @@ struct mtmd_gen_out {
     // for MTMD_GEN_PROCESS_TYPE_GEN_WAV
     const float * audio;
     size_t        n_samples;
-    const char * state_data;
-    size_t       state_size;
 };
 
+struct mtmd_gen_state_params {
+    uint32_t seed; // UINT32_MAX for random
+};
+
+MTMD_API struct mtmd_gen_state_params mtmd_gen_state_params_default(void);
+
+// returns nullptr if the model does not support audio generation
+MTMD_API mtmd_gen_state * mtmd_gen_state_init(mtmd_context * ctx, struct mtmd_gen_state_params params);
+MTMD_API void             mtmd_gen_state_free(mtmd_gen_state * state);
+
+// clear the decoder state (cold start), keep the RNG stream
+MTMD_API void             mtmd_gen_state_reset(mtmd_gen_state * state);
+
 // defaults tuned for the loaded pipeline, callers override only what they care about
 MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
 
-// note: this API is stateless, caller must handle state management and audio frame accumulation
+// run one step and update the state in place
+// the caller drives the backbone and collects the audio frames
 MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
+                                mtmd_gen_state * state,
                                 const struct mtmd_gen_inp * inp,
                                 struct mtmd_gen_out * out);
 
@@ -544,6 +557,11 @@ struct mtmd_audio_stream_deleter {
 };
 using audio_stream_ptr = std::unique_ptr<mtmd_audio_stream, mtmd_audio_stream_deleter>;
 
+struct mtmd_gen_state_deleter {
+    void operator()(mtmd_gen_state * val) { mtmd_gen_state_free(val); }
+};
+using gen_state_ptr = std::unique_ptr<mtmd_gen_state, mtmd_gen_state_deleter>;
+
 struct bitmap {
     bitmap_ptr ptr;
     bitmap() : ptr(nullptr) {}

tools/server/README.md

@@ -1719,6 +1719,8 @@ See [Anthropic Messages API documentation](https://docs.anthropic.com/en/api/mes
 
 `top_k`: Top-k sampling
 
+`extra_params`: llama.cpp parameters that the Anthropic API has no field for, such as `seed`, `min_p` or `repeat_penalty`, as an object (e.g. `{"seed": 42, "min_p": 0.05}`). Its keys are merged into the converted request and behave as on `/v1/chat/completions`.
+
 `stop_sequences`: Array of stop sequences
 
 `stream`: Enable streaming (default: false)
@@ -1739,7 +1741,8 @@ curl http://localhost:9931/v1/messages \
     "system": "You are a helpful assistant.",
     "messages": [
       {"role": "user", "content": "Hello!"}
-    ]
+    ],
+    "extra_params": {"seed": 42, "min_p": 0.05}
   }'
 ```
 

tools/server/server-chat.cpp

@@ -615,6 +615,13 @@ json server_chat_convert_anthropic_to_oai(const json & body) {
         }
     }
 
+    // llama.cpp params (e.g. seed, min_p) go under extra_params and are merged into the converted body
+    if (body.contains("extra_params") && body.at("extra_params").is_object()) {
+        for (const auto & [key, value] : body.at("extra_params").items()) {
+            oai_body[key] = value;
+        }
+    }
+
     return oai_body;
 }
 

vendor/cpp-httplib/CMakeLists.txt

@@ -43,7 +43,7 @@ if (LLAMA_BUILD_BORINGSSL)
     set(FIPS OFF CACHE BOOL "Enable FIPS (BoringSSL)")
 
     set(BORINGSSL_GIT "https://boringssl.googlesource.com/boringssl" CACHE STRING "BoringSSL git repository")
-    set(BORINGSSL_VERSION "0.20260929.0" CACHE STRING "BoringSSL version")
+    set(BORINGSSL_VERSION "0.20261020.0" CACHE STRING "BoringSSL version")
 
     message(STATUS "Fetching BoringSSL version ${BORINGSSL_VERSION}")