Merge remote-tracking branch 'origin/main' into ik/better_indexer_mask

examples/quantize/quantize.cpp

@@ -396,6 +396,12 @@ int main(int argc, char ** argv) {
             } else {
                 usage(argv[0]);
             }
+        } else if (strcmp(argv[arg_idx], "--per-layer-token-embedding-type") == 0) {
+            if (arg_idx < argc-1) {
+                params.per_layer_token_embedding_type = parse_ggml_type(argv[++arg_idx]);
+            } else {
+                usage(argv[0]);
+            }
         } else if (strcmp(argv[arg_idx], "--ffn-gate-inp-type") == 0) {
             if (arg_idx < argc-1) {
                 params.ffn_gate_inp_type = parse_ggml_type(argv[++arg_idx]);

ggml/CMakeLists.txt

@@ -82,7 +82,7 @@ option(GGML_CPU_HBM     "ggml: use memkind for CPU HBM" OFF)
 
 option(GGML_AVX         "ggml: enable AVX"              ${INS_ENB})
 option(GGML_AVX2        "ggml: enable AVX2"             ${INS_ENB})
-option(GGML_AVXVNNI     "ggml: enable AVX-VNNI"         ${INS_ENB})
+option(GGML_AVXVNNI     "ggml: enable AVX-VNNI"         OFF)
 option(GGML_AVX512      "ggml: enable AVX512"           OFF)
 option(GGML_AVX512_VBMI "ggml: enable AVX512-VBMI"      OFF)
 option(GGML_AVX512_VNNI "ggml: enable AVX512-VNNI"      OFF)

ggml/include/ggml.h

@@ -706,6 +706,7 @@ extern "C" {
         GGML_OP_BLEND,
         GGML_OP_INDEXER_TOPK,
         GGML_OP_MASK_TOPK,
+        GGML_OP_SINKHORN,
 
         GGML_OP_COUNT,
     };
@@ -2584,6 +2585,22 @@ extern "C" {
             enum ggml_unary_op    op,
             int                   n_top_k);
 
+    // Sinkhorn normalization of a flat [S*S, T] batch of S x S matrices into
+    // doubly-stochastic form: softmax over columns, then column normalization,
+    // then (n_iters - 1) rounds of row + column normalization (ends on columns).
+    // The flat input is row-major (column index fastest). eps, when non-zero, is
+    // added to the softmax output and to every normalization sum before dividing.
+    // With output_transposed the result is [S, S, T] with ne0 = row, ne1 = column
+    // (ready for out[c] = sum_r m[r,c] * residual[r] consumers); otherwise the
+    // bare input layout (ne0 = column) is kept.
+    GGML_API struct ggml_tensor * ggml_sinkhorn(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   S,
+            int                   n_iters,
+            float                 eps,
+            bool                  output_transposed);
+
     // custom operators
 
     typedef void (*ggml_unary_op_f32_t) (const int, float *, const float *);

ggml/src/ggml-cuda.cu

@@ -56,6 +56,7 @@
 #include "ggml-cuda/reduce.cuh"
 #include "ggml-cuda/tri.cuh"
 #include "ggml-cuda/delta-net.cuh"
+#include "ggml-cuda/sinkhorn.cuh"
 #include "ggml-cuda/blend.cuh"
 #include "ggml-cuda/indexer_topk.cuh"
 
@@ -4127,6 +4128,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_DELTA_NET:
             ggml_cuda_op_delta_net(ctx, dst);
             break;
+        case GGML_OP_SINKHORN:
+            ggml_cuda_op_sinkhorn(ctx, dst);
+            break;
         case GGML_OP_FLASH_ATTN_EXT:
             ggml_cuda_flash_attn_ext(ctx, dst);
             break;
@@ -5042,6 +5046,11 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
         case GGML_OP_INDEXER_TOPK:
         case GGML_OP_MASK_TOPK:
             return true;
+        case GGML_OP_SINKHORN: {
+            const int sink_s = op->op_params[0];
+            return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 &&
+                   sink_s >= 1 && sink_s <= 8 && op->src[0]->ne[0] == (int64_t) sink_s*sink_s;
+        }
         case GGML_OP_FLASH_ATTN_EXT:
 #if defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__)
             return (op->src[0]->ne[0] == 64 && op->src[1]->type == GGML_TYPE_F16) || op->src[0]->ne[0] == 128;

ggml/src/ggml-cuda/sinkhorn.cu

@@ -0,0 +1,103 @@
+#include "common.cuh"
+#include "sinkhorn.cuh"
+
+// Sinkhorn normalization of T independent S x S matrices (S <= 8, so a matrix is
+// at most 64 floats). One thread per token: the whole matrix lives in a thread-local
+// array and the 6-node-per-iteration graph chain collapses into this single kernel.
+// Semantics match the reference: softmax over columns, column normalization,
+// then (iters - 1) rounds of row + column normalization (ends on columns).
+// Input is the flat [S*S, T] row-major tensor (column index fastest); output is
+// [S, S, T] with ne0 = row, i.e. transposed on write.
+
+template <int S>
+static __global__ void k_sinkhorn(const float * __restrict__ x, float * __restrict__ dst,
+                                  const int64_t T, const int iters, const float eps,
+                                  const int transposed, const int64_t nb1) {
+    const int64_t t = (int64_t) blockIdx.x*blockDim.x + threadIdx.x;
+    if (t >= T) {
+        return;
+    }
+
+    const float * xt = (const float *)((const char *) x + t*nb1);
+    float m[S*S];
+
+    #pragma unroll
+    for (int r = 0; r < S; ++r) {
+        float mx = xt[r*S];
+        for (int c = 1; c < S; ++c) mx = fmaxf(mx, xt[r*S + c]);
+        float sum = 0.0f;
+        for (int c = 0; c < S; ++c) { m[r*S + c] = expf(xt[r*S + c] - mx); sum += m[r*S + c]; }
+        for (int c = 0; c < S; ++c) m[r*S + c] = m[r*S + c]/sum + eps;
+    }
+    #pragma unroll
+    for (int c = 0; c < S; ++c) {
+        float sum = eps;
+        for (int r = 0; r < S; ++r) sum += m[r*S + c];
+        for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+    }
+    for (int i = 0; i < iters - 1; ++i) {
+        #pragma unroll
+        for (int r = 0; r < S; ++r) {
+            float sum = eps;
+            for (int c = 0; c < S; ++c) sum += m[r*S + c];
+            for (int c = 0; c < S; ++c) m[r*S + c] /= sum;
+        }
+        #pragma unroll
+        for (int c = 0; c < S; ++c) {
+            float sum = eps;
+            for (int r = 0; r < S; ++r) sum += m[r*S + c];
+            for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+        }
+    }
+
+    float * yt = dst + t*S*S;
+    if (transposed) {
+        #pragma unroll
+        for (int c = 0; c < S; ++c) {
+            for (int r = 0; r < S; ++r) yt[c*S + r] = m[r*S + c];
+        }
+    } else {
+        #pragma unroll
+        for (int k = 0; k < S*S; ++k) yt[k] = m[k];
+    }
+}
+
+void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    const int S     = dst->op_params[0];
+    const int iters = dst->op_params[1];
+    float eps;
+    memcpy(&eps, &dst->op_params[2], sizeof(float));
+    const int transposed = dst->op_params[3];
+    const int64_t T = src0->ne[1];
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(S >= 1 && S <= 8);
+    GGML_ASSERT(src0->ne[0] == (int64_t) S * S);
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    if (T == 0) {
+        return;
+    }
+
+    const int block = 256;
+    const int64_t grid = (T + block - 1)/block;
+    cudaStream_t stream = ctx.stream();
+
+    const float * x = (const float *) src0->data;
+    float * y = (float *) dst->data;
+
+    switch (S) {
+        case 1: k_sinkhorn<1><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 2: k_sinkhorn<2><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 3: k_sinkhorn<3><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 4: k_sinkhorn<4><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 5: k_sinkhorn<5><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 6: k_sinkhorn<6><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 7: k_sinkhorn<7><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        case 8: k_sinkhorn<8><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
+        default: GGML_ABORT("sinkhorn: unsupported S");
+    }
+}

ggml/src/ggml-cuda/sinkhorn.cuh

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

ggml/src/ggml.c

@@ -4334,9 +4334,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "BLEND",
     "INDEXER_TOPK",
     "MASK_TOPK",
+    "SINKHORN",
 };
 
-static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105");
+static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -4457,10 +4458,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "blend(a,b,c)",
     "indexer_topk(k, q, w, mask)",
     "mask_topk(mask, topk)",
+    "sinkhorn(x)",
 
 };
 
-static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105");
+static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -10141,6 +10143,32 @@ struct ggml_tensor * ggml_indexer_topk(
 }
 
 
+struct ggml_tensor * ggml_sinkhorn(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,
+            int                   S,
+            int                   n_iters,
+            float                 eps,
+            bool                  output_transposed) {
+    GGML_ASSERT(eps >= 0.0f);
+    GGML_ASSERT(a->type == GGML_TYPE_F32);
+    GGML_ASSERT(S >= 1 && S <= 8);
+    GGML_ASSERT(a->ne[0] == (int64_t) S * S);
+    GGML_ASSERT(a->ne[2] == 1 && a->ne[3] == 1);
+    GGML_ASSERT(n_iters >= 1);
+    GGML_ASSERT(ggml_is_contiguous(a));
+
+    struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, S, a->ne[1]);
+    result->op = GGML_OP_SINKHORN;
+    result->op_params[0] = S;
+    result->op_params[1] = n_iters;
+    memcpy(&result->op_params[2], &eps, sizeof(float));
+    result->op_params[3] = output_transposed ? 1 : 0;
+    result->src[0] = a;
+
+    return result;
+}
+
 // ggml_fill
 
 static struct ggml_tensor * ggml_fill_impl(
@@ -23111,6 +23139,83 @@ static void ggml_compute_forward_delta_net(
     }
 }
 
+// ggml_compute_forward_sinkhorn
+
+static void ggml_compute_forward_sinkhorn_f32(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0];
+
+    const int S     = dst->op_params[0];
+    const int iters = dst->op_params[1];
+    float eps;
+    memcpy(&eps, &dst->op_params[2], sizeof(float));
+    const int transposed = dst->op_params[3];
+    const int64_t T = src0->ne[1];
+
+    GGML_ASSERT(S >= 1 && S <= 8);
+    GGML_ASSERT(src0->ne[0] == (int64_t) S * S);
+    GGML_ASSERT(iters >= 1);
+
+    // one token is S*S floats (16 at S=4): parallelize over tokens only
+    const int64_t t0 = (T *  params->ith   ) / params->nth;
+    const int64_t t1 = (T * (params->ith+1)) / params->nth;
+
+    float m[64];
+
+    for (int64_t t = t0; t < t1; ++t) {
+        const float * x = (const float *)((const char *)src0->data + t*src0->nb[1]);
+        float       * y = (float       *)((      char *)dst->data  + t*dst->nb[2]);
+
+        // softmax over columns c for each row r; flat input is row-major (c fastest)
+        for (int r = 0; r < S; ++r) {
+            float mx = x[r*S];
+            for (int c = 1; c < S; ++c) mx = MAX(mx, x[r*S + c]);
+            float sum = 0.0f;
+            for (int c = 0; c < S; ++c) { m[r*S + c] = expf(x[r*S + c] - mx); sum += m[r*S + c]; }
+            for (int c = 0; c < S; ++c) m[r*S + c] = m[r*S + c]/sum + eps;
+        }
+        // column normalization first, then (iters - 1) rounds of row + column: ends on columns
+        for (int c = 0; c < S; ++c) {
+            float sum = eps;
+            for (int r = 0; r < S; ++r) sum += m[r*S + c];
+            for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+        }
+        for (int i = 0; i < iters - 1; ++i) {
+            for (int r = 0; r < S; ++r) {
+                float sum = eps;
+                for (int c = 0; c < S; ++c) sum += m[r*S + c];
+                for (int c = 0; c < S; ++c) m[r*S + c] /= sum;
+            }
+            for (int c = 0; c < S; ++c) {
+                float sum = eps;
+                for (int r = 0; r < S; ++r) sum += m[r*S + c];
+                for (int r = 0; r < S; ++r) m[r*S + c] /= sum;
+            }
+        }
+        if (transposed) {
+            // dst is [row, col, T] (ne0 = row): transpose on write
+            for (int c = 0; c < S; ++c) {
+                for (int r = 0; r < S; ++r) y[c*S + r] = m[r*S + c];
+            }
+        } else {
+            for (int k = 0; k < S*S; ++k) y[k] = m[k];
+        }
+    }
+}
+
+static void ggml_compute_forward_sinkhorn(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    switch (dst->src[0]->type) {
+        case GGML_TYPE_F32:
+            ggml_compute_forward_sinkhorn_f32(params, dst);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
 // ggml_compute_forward_win_part
 
 static void ggml_compute_forward_win_part_f32(
@@ -24853,6 +24958,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml
             {
                 ggml_compute_forward_delta_net(params, tensor);
             } break;
+        case GGML_OP_SINKHORN:
+            {
+                ggml_compute_forward_sinkhorn(params, tensor);
+            } break;
         case GGML_OP_INDEXER_TOPK:
             {
                 if (!iqk_indexer_topk(tensor, params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, params->ith, params->nth)) {
@@ -25926,6 +26035,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
         case GGML_OP_DELTA_NET:
         case GGML_OP_INDEXER_TOPK:
         case GGML_OP_MASK_TOPK:
+        case GGML_OP_SINKHORN:
             {
                 GGML_ABORT("fatal error"); // TODO: not implemented
             }
@@ -26664,6 +26774,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_DELTA_NET:
         case GGML_OP_INDEXER_TOPK:
         case GGML_OP_MASK_TOPK:
+        case GGML_OP_SINKHORN:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/iqk/iqk_quantize.cpp

@@ -8353,10 +8353,7 @@ const Modify * get_modify_info(ggml_type type) {
     return it != k_mod_map.end() ? &it->second : nullptr;
 }
 bool is_forbidden_tensor(const std::string& name) {
-    static const std::string kTokenEmbd{"token_embd.weight"};
-    if (name == kTokenEmbd) return true;
-    //if (auto pos = name.find("attn_kv_b.weight"); pos != std::string::npos) return true;
-    return false;
+    return (name == "token_embd.weight" || name == "per_layer_token_embd.weight");
 }
 }
 

include/llama.h

@@ -523,6 +523,7 @@ extern "C" {
         enum llama_ftype ftype;              // quantize to this llama_ftype
         enum ggml_type output_tensor_type;   // output tensor type
         enum ggml_type token_embedding_type; // token embeddings tensor type
+        enum ggml_type per_layer_token_embedding_type; // token embeddings tensor type
         enum ggml_type attn_q_type;          // attention query tensor type
         enum ggml_type attn_k_type;          // attention key tensor type
         enum ggml_type attn_v_type;          // attention value tensor type

src/graphs/build_openpangu.cpp

@@ -222,39 +222,6 @@ static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf,
     return out;
 }
 
-// --- mHC Sinkhorn: h_res [S*S, T] -> doubly-stochastic per token, 20 iters (ends on col norm) ---
-static ggml_tensor * openpangu_sinkhorn(ggml_context * ctx, ggml_tensor * h_res_flat,
-                                        int64_t S, int64_t T, int iters, float hc_eps) {
-    // The flat h_res is torch [r,c] row-major (c fastest), so a bare reshape gives ne0=col.
-    // Transpose once so ne0=row(S), ne1=col(S): every axis op below then matches the
-    // reference _mhc_sinkhorn_naive (softmax over col, first norm over row, end on col-sum=1)
-    // and mhc_post's out[c] = sum_r m[r,c]*residual[r].
-    (void) hc_eps; // softmax outputs are strictly positive, so the eps is numerically inert here
-    ggml_tensor * m = ggml_reshape_3d(ctx, h_res_flat, S, S, T);   // ne0=col (bare reshape)
-    // ref softmaxes h_res over columns; a bare reshape already has ne0=col, so soft_max
-    // (over ne0) hits the column axis directly -- no pre-permute round-trip needed.
-    m = ggml_soft_max(ctx, m);                                 // softmax over col
-    m = ggml_cont(ctx, ggml_permute(ctx, m, 1, 0, 2, 3));      // transpose once -> [row, col, T]
-
-    auto col_norm = [&](ggml_tensor * a) {
-        ggml_tensor * col_sum = ggml_sum_rows(ctx, a);          // sums ne0(row) -> [1, col, T]
-        return ggml_div(ctx, a, col_sum);                       // broadcast [1,col,T] over rows
-    };
-    auto row_norm = [&](ggml_tensor * a) {
-        ggml_tensor * ap = ggml_cont(ctx, ggml_permute(ctx, a, 1, 0, 2, 3)); // [col,row,T]
-        ggml_tensor * row_sum = ggml_sum_rows(ctx, ap);         // [1, row, T]
-        ggml_tensor * out = ggml_div(ctx, ap, row_sum);
-        return ggml_cont(ctx, ggml_permute(ctx, out, 1, 0, 2, 3)); // back [row,col,T]
-    };
-
-    m = col_norm(m);
-    for (int i = 0; i < iters - 1; ++i) {
-        m = row_norm(m);
-        m = col_norm(m);
-    }
-    return m; // [row(S), col(S), T]
-}
-
 // Attention sublayer body, shared by the base layers and the NextN/MTP head.
 // x_normed = input-layernormed hidden [n_embd, T]; returns post-o_proj output [n_embd, T].
 // conv_state is the recurrent MoME state slot. seq_qnext is the [1, T] sequence-id input
@@ -876,7 +843,6 @@ ggml_cgraph * llm_build_context::build_openpangu() {
     const int64_t n_embd_head_k = hparams.n_embd_head_k(0);                // 192
     const int64_t S             = hparams.mhc_num_stream;                  // 4
     const int    sink_iters     = (int) hparams.mhc_recur_norm;            // 20
-    const float  hc_eps         = 1e-6f;
     const float  kq_scale       = 1.0f / sqrtf(float(n_embd_head_k));
 
 
@@ -1076,7 +1042,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
         // cont is required: the CUDA broadcast-mul path misreads strided views (h_pre is a
         // row-slice of mixes), while CPU handles the strides — token 0 right, tokens 1+ garbage
         h_pre = ggml_add(ctx0, ggml_mul(ctx0, ggml_cont(ctx0, h_pre), a_pre), b_pre);  // broadcast scalar + [S]
-        h_pre = ggml_sigmoid(ctx0, h_pre);                            // [S,T] (+hc_eps omitted, inert)
+        h_pre = ggml_sigmoid(ctx0, h_pre);                            // [S,T] (+eps omitted, inert)
 
         // combine: x[h,t] = sum_s h_pre[s,t] * R[h,s,t]
         ggml_tensor * hpre3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_pre), 1, S, n_tokens);
@@ -1101,7 +1067,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
         h_post = ggml_scale(ctx0, ggml_sigmoid(ctx0, h_post), 2.0f);  // 2*sigmoid, [S,T]
 
         ggml_tensor * m = ggml_add(ctx0, ggml_mul(ctx0, h_res, a_res), b_res); // [S*S,T]
-        m = openpangu_sinkhorn(ctx0, m, S, n_tokens, sink_iters, hc_eps);      // [row S, col S, T]
+        m = ggml_sinkhorn(ctx0, m, (int) S, sink_iters, 0.0f, /*output_transposed=*/true); // [row S, col S, T]
 
         // term1: h_post[s,t]*y[h,t] -> [H,S,T]
         ggml_tensor * y3 = ggml_reshape_3d(ctx0, y, n_embd, 1, n_tokens);

src/llama-quantize.cpp

@@ -409,6 +409,29 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
                 new_type = GGML_TYPE_IQ4_NL;
             }
         }
+    } else if (name == "per_layer_token_embd.weight") {
+        if (qs.params->per_layer_token_embedding_type < GGML_TYPE_COUNT) {
+            new_type = qs.params->per_layer_token_embedding_type;
+        } else {
+            if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ2_XS ||
+                ftype == LLAMA_FTYPE_MOSTLY_IQ1_S   || ftype == LLAMA_FTYPE_MOSTLY_IQ1_M  ||
+                ftype == LLAMA_FTYPE_MOSTLY_IQ2_XXS_R4 || ftype == LLAMA_FTYPE_MOSTLY_IQ2_XS_R4 ||
+                ftype == LLAMA_FTYPE_MOSTLY_IQ1_S_R4 || ftype == LLAMA_FTYPE_MOSTLY_IQ1_M_R4) {
+                new_type = GGML_TYPE_Q2_K;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ2_S || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M || ftype == LLAMA_FTYPE_MOSTLY_IQ2_M_R4) {
+                new_type = GGML_TYPE_IQ3_S;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS || ftype == LLAMA_FTYPE_MOSTLY_IQ3_KT) {
+                new_type = GGML_TYPE_IQ3_S;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS_R4) {
+                new_type = GGML_TYPE_IQ3_K;
+            }
+            else if (ftype == LLAMA_FTYPE_MOSTLY_IQ1_BN || ftype == LLAMA_FTYPE_MOSTLY_IQ2_BN || ftype == LLAMA_FTYPE_MOSTLY_IQ2_BN_R4) {
+                new_type = GGML_TYPE_IQ4_NL;
+            }
+        }
     } else if (ftype == LLAMA_FTYPE_MOSTLY_IQ1_S_R4 || ftype == LLAMA_FTYPE_MOSTLY_IQ1_M_R4) {
         if (name.find("attn_v.weight") != std::string::npos) {
             if (qs.model.hparams.n_expert >= 4 || qs.model.hparams.n_gqa() >= 4) new_type = GGML_TYPE_IQ4_K_R4;
@@ -830,6 +853,15 @@ static ggml_type llama_tensor_get_type(quantize_state_internal & qs, ggml_type n
         }
     }
 
+    if (name == "per_layer_token_embd.weight") {
+        auto working_type = interleaved_properties(new_type).first;
+        if (working_type != new_type) {
+            printf("\n============ Per-layer Token embeddings cannot be quantized with row-interleaved quants\n");
+            printf("---> Changed %s to %s\n", ggml_type_name(new_type), ggml_type_name(working_type));
+            new_type = working_type;
+        }
+    }
+
     return new_type;
 }
 
@@ -1522,6 +1554,9 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
             if (params->token_embedding_type < GGML_TYPE_COUNT && strcmp(tensor->name, "token_embd.weight") == 0) {
                 new_type = params->token_embedding_type;
             }
+            if (params->per_layer_token_embedding_type < GGML_TYPE_COUNT && strcmp(tensor->name, "per_layer_token_embd.weight") == 0) {
+                new_type = params->per_layer_token_embedding_type;
+            }
             if (params->output_tensor_type < GGML_TYPE_COUNT && strcmp(tensor->name, "output.weight") == 0) {
                 new_type = params->output_tensor_type;
             }

src/llama.cpp

@@ -7186,34 +7186,35 @@ struct llama_context_params llama_context_default_params() {
 
 struct llama_model_quantize_params llama_model_quantize_default_params() {
     struct llama_model_quantize_params result = {
-        /*.nthread                     =*/ 0,
-        /*.ftype                       =*/ LLAMA_FTYPE_MOSTLY_Q5_1,
-        /*.output_tensor_type          =*/ GGML_TYPE_COUNT,
-        /*.token_embedding_type        =*/ GGML_TYPE_COUNT,
-        /*.attn_q_type                 =*/ GGML_TYPE_COUNT,
-        /*.attn_k_type                 =*/ GGML_TYPE_COUNT,
-        /*.attn_v_type                 =*/ GGML_TYPE_COUNT,
-        /*.attn_qkv_type               =*/ GGML_TYPE_COUNT,
-        /*.attn_output_type            =*/ GGML_TYPE_COUNT,
-        /*.ffn_gate_type               =*/ GGML_TYPE_COUNT,
-        /*.ffn_down_type               =*/ GGML_TYPE_COUNT,
-        /*.ffn_up_type                 =*/ GGML_TYPE_COUNT,
-        /*.ffn_gat_inp_type            =*/ GGML_TYPE_COUNT,
-        /*.extra_output_type           =*/ GGML_TYPE_COUNT,
-        /*.allow_requantize            =*/ false,
-        /*.quantize_output_tensor      =*/ true,
-        /*.only_copy                   =*/ false,
-        /*.pure                        =*/ false,
-        /*.keep_split                  =*/ false,
-        /*.ignore_imatrix_rules        =*/ false,
-        /*.only_repack                 =*/ false,
-        /*.dry_run                     =*/ false,
-        /*.partial_requant             =*/ false,
-        /*.imatrix                     =*/ nullptr,
-        /*.kv_overrides                =*/ nullptr,
-        /*.custom_quants               =*/ nullptr,
-        /*.repack_pattern              =*/ nullptr,
-        /*.user_data                   =*/ nullptr,
+        /*.nthread                        =*/ 0,
+        /*.ftype                          =*/ LLAMA_FTYPE_MOSTLY_Q5_1,
+        /*.output_tensor_type             =*/ GGML_TYPE_COUNT,
+        /*.token_embedding_type           =*/ GGML_TYPE_COUNT,
+        /*.per_layer_token_embedding_type =*/ GGML_TYPE_COUNT,
+        /*.attn_q_type                    =*/ GGML_TYPE_COUNT,
+        /*.attn_k_type                    =*/ GGML_TYPE_COUNT,
+        /*.attn_v_type                    =*/ GGML_TYPE_COUNT,
+        /*.attn_qkv_type                  =*/ GGML_TYPE_COUNT,
+        /*.attn_output_type               =*/ GGML_TYPE_COUNT,
+        /*.ffn_gate_type                  =*/ GGML_TYPE_COUNT,
+        /*.ffn_down_type                  =*/ GGML_TYPE_COUNT,
+        /*.ffn_up_type                    =*/ GGML_TYPE_COUNT,
+        /*.ffn_gat_inp_type               =*/ GGML_TYPE_COUNT,
+        /*.extra_output_type              =*/ GGML_TYPE_COUNT,
+        /*.allow_requantize               =*/ false,
+        /*.quantize_output_tensor         =*/ true,
+        /*.only_copy                      =*/ false,
+        /*.pure                           =*/ false,
+        /*.keep_split                     =*/ false,
+        /*.ignore_imatrix_rules           =*/ false,
+        /*.only_repack                    =*/ false,
+        /*.dry_run                        =*/ false,
+        /*.partial_requant                =*/ false,
+        /*.imatrix                        =*/ nullptr,
+        /*.kv_overrides                   =*/ nullptr,
+        /*.custom_quants                  =*/ nullptr,
+        /*.repack_pattern                 =*/ nullptr,
+        /*.user_data                      =*/ nullptr,
     };
 
     return result;