model: Ling-3.0 (bailingmoe3) runtime support (#2295)

* model: Ling-3.0 (bailingmoe3) runtime support

* model: Ling-3.0-tiny support

---------

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>

ggml/src/CMakeLists.txt

@@ -180,6 +180,7 @@ if (GGML_IQK_MUL_MAT)
     message(STATUS "Using optimized iqk matrix multiplications")
     add_compile_definitions(GGML_USE_IQK_MULMAT)
     set(GGML_SOURCES_IQK_MM iqk/iqk_mul_mat.cpp
+                            iqk/iqk_kda.cpp
                             iqk/iqk_flash_attn.cpp
                             iqk/fa/iqk_fa_576_512.cpp
                             iqk/fa/iqk_fa_512_512.cpp

ggml/src/ggml-cuda.cu

@@ -58,6 +58,7 @@
 #include "ggml-cuda/reduce.cuh"
 #include "ggml-cuda/tri.cuh"
 #include "ggml-cuda/delta-net.cuh"
+#include "ggml-cuda/kda.cuh"
 #include "ggml-cuda/sinkhorn.cuh"
 #include "ggml-cuda/latent_attn.cuh"
 #include "ggml-cuda/blend.cuh"
@@ -4140,11 +4141,12 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
             ggml_cuda_op_solve_tri(ctx, dst);
             break;
         case GGML_OP_DELTA_NET: {
+            const auto op_delta_net = dst->src[3]->ne[1] == 1 ? ggml_cuda_op_delta_net : ggml_cuda_op_kda;
             const int j = fusion ? ggml_delta_net_find_state_cpy(cgraph, i) : -1;
             if (j >= 0) {
                 ggml_tensor fused = *dst;
                 fused.src[7] = cgraph->nodes[j]->src[1];
-                ggml_cuda_op_delta_net(ctx, &fused);
+                op_delta_net(ctx, &fused);
 #ifdef USE_CUDA_GRAPH
                 // claim the entry of the copy that is not going to be launched
                 if (ctx.cur_graph && ctx.cur_graph->use_cpy_indirection) {
@@ -4153,7 +4155,7 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
 #endif
                 i = j;
             } else {
-                ggml_cuda_op_delta_net(ctx, dst);
+                op_delta_net(ctx, dst);
             }
         } break;
         case GGML_OP_SINKHORN:

ggml/src/ggml-cuda/kda.cu

@@ -0,0 +1,268 @@
+#include "common.cuh"
+#include "kda.cuh"
+
+#include <cstdint>
+
+__device__ __forceinline__ float sigmoid_f(float x) {
+    return 1.0f / (1.0f + expf(-x));
+}
+
+template <int block_size>
+__device__ __forceinline__ float reduce_sum(float x, float * s) {
+    x = warp_reduce_sum(x);
+    if constexpr (block_size > WARP_SIZE) {
+        //__shared__ float s[block_size/WARP_SIZE];
+        int warp_id = threadIdx.x / WARP_SIZE;
+        int lane_id = threadIdx.x % WARP_SIZE;
+        if (lane_id == 0) {
+            s[warp_id] = x;
+        }
+        __syncthreads();
+        x = lane_id < block_size/WARP_SIZE ? s[lane_id] : 0.0f;
+        x = warp_reduce_sum(x);
+    }
+    return x;
+}
+
+template <int HEAD_DIM, int block_size>
+__global__ void kda_recurrent_f32(
+    const float * __restrict__ q,         // [HEAD_DIM, n_tokens, n_heads, n_seqs]
+    const float * __restrict__ k,         // [HEAD_DIM, n_tokens, n_heads, n_seqs]
+    const float * __restrict__ v,         // [HEAD_DIM, n_tokens, n_heads, n_seqs]
+    const float * __restrict__ g,
+    const float * __restrict__ beta_in,   // [1, n_tokens, n_heads, n_seqs]
+    const float * state_in,               // [HEAD_DIM, HEAD_DIM*n_heads, 1, n_seqs], aliases state_out when fused
+    float * __restrict__ dst,             // output
+    float * state_out,                    // new state
+    float * __restrict__ saved_states,
+    const int32_t n_heads,
+    const int32_t gqa_ratio,
+    const int32_t repeat_type,
+    const int32_t n_tokens,
+    const int32_t n_seqs,
+    size_t vnb1, size_t vnb2, size_t vnb3,
+    size_t gnb0, size_t gnb1, size_t gnb2, size_t gnb3,
+    size_t bnb1, size_t bnb2, size_t bnb3) {
+    constexpr int32_t warps_per_head = HEAD_DIM/WARP_SIZE;
+    const int32_t batch_idx = blockIdx.x / (warps_per_head*n_heads);
+    const int32_t sub_head_idx = blockIdx.x % (warps_per_head*n_heads);
+    const int32_t head_idx = sub_head_idx / warps_per_head;
+    const int32_t sub_idx = sub_head_idx % warps_per_head;
+    const int32_t head_idx_kq = repeat_type == 0 ? head_idx / gqa_ratio : head_idx % (n_heads/gqa_ratio);
+    const int32_t tid = threadIdx.x;
+
+    // Strides for input tensors (column-major)
+    // Q/K/V: [HEAD_DIM, n_tokens, n_heads, n_seqs]
+    const int32_t qkv_stride_token = HEAD_DIM;
+    const int32_t qkv_stride_head = HEAD_DIM * n_tokens;
+    const int32_t qkv_stride_batch = HEAD_DIM * n_tokens * n_heads;
+    const int32_t qkv_stride_batch_kq = qkv_stride_batch / gqa_ratio;
+
+    // State: [HEAD_DIM, HEAD_DIM*n_heads, 1, n_seqs]
+    // For head h: columns h*HEAD_DIM to (h+1)*HEAD_DIM
+    // state[row, col] for head h = state[row, h*HEAD_DIM + col]
+    // Linear index: row + (h*HEAD_DIM + col) * HEAD_DIM = row + h*HEAD_DIM^2 + col*HEAD_DIM
+    const int32_t state_head_offset = head_idx * HEAD_DIM * HEAD_DIM;
+    const int32_t state_batch_stride = HEAD_DIM * HEAD_DIM * n_heads;
+
+    // State step stride for save_all_states: HEAD_DIM^2 * n_heads * n_seqs
+    const int32_t state_step_stride = HEAD_DIM * HEAD_DIM * n_heads * n_seqs;
+
+    // Pointers for this batch/head
+    const float * q_ptr = q + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
+    const float * k_ptr = k + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
+    const float * v_ptr = v + batch_idx * vnb3 + head_idx * vnb2;
+    const float * g_ptr = g + batch_idx * gnb3 + head_idx * gnb2;
+    const float * beta_ptr = beta_in + batch_idx * bnb3 + head_idx * bnb2;
+    const float * state_src = state_in + batch_idx * state_batch_stride + state_head_offset;
+
+    // Output layout: [head_v_dim, num_v_heads, n_seq_tokens, n_seqs]
+    // For [dim, head, token, batch]: index = dim + head*S_v + token*S_v*H_v + batch*S_v*H_v*n_tokens
+    float * out_base = dst + batch_idx * (HEAD_DIM * n_heads * n_tokens) + head_idx * HEAD_DIM;
+    const int32_t out_token_stride = HEAD_DIM * n_heads;  // stride between tokens
+    float * state_dst = state_out + batch_idx * state_batch_stride + state_head_offset;
+
+    // Shared memory for current token's Q, K, V (normalized), and intermediate results
+    extern __shared__ float smem[];
+    float * sQ = smem;                          // HEAD_DIM
+    float * sK = sQ + HEAD_DIM;                 // HEAD_DIM
+
+    const float scale = rsqrtf((float)HEAD_DIM);
+
+    __shared__ float sum_helper[block_size/WARP_SIZE];
+
+    constexpr int32_t num_warps = block_size/WARP_SIZE;
+    const int32_t row = tid % WARP_SIZE;
+    const int32_t col_idx_0 = tid / WARP_SIZE;
+    const int32_t row_out = row + sub_idx * WARP_SIZE;
+
+    // Keep the state in registers, copy the final state to its destination at the end
+    float state_local[HEAD_DIM/num_warps];
+    for (int32_t i = 0; i < HEAD_DIM/num_warps; ++i) {
+        int32_t col = num_warps*i + col_idx_0;
+        state_local[i] = state_src[col*HEAD_DIM + row_out];
+    }
+
+    constexpr int32_t WARP_SIZE_S = WARP_SIZE + 1;
+    constexpr int32_t num_stored_rows = block_size/WARP_SIZE;
+    __shared__ float all_sum[2*WARP_SIZE_S*num_stored_rows];
+    auto all_sum1 = all_sum;
+    auto all_sum2 = all_sum1 + WARP_SIZE_S*num_stored_rows;
+
+    for (int32_t t = 0; t < n_tokens; t++) {
+        float sum_kq = 0.0f;
+        for (int32_t i = tid; i < HEAD_DIM; i += block_size) {
+            sQ[i] = q_ptr[t * qkv_stride_token + i] * scale;
+            sK[i] = k_ptr[t * qkv_stride_token + i];
+            sum_kq += sK[i] * sQ[i];
+        }
+
+        float attn_score = reduce_sum<block_size>(sum_kq, sum_helper);
+
+        float beta_val = sigmoid_f(beta_ptr[t*bnb1]);
+
+        float sum1 = 0, sum2 = 0;
+#pragma unroll
+        for (int32_t i = 0; i < HEAD_DIM/num_warps; ++i) {
+            int32_t col = num_warps*i + col_idx_0;
+            float decay = row == 0 ? expf(fminf(g_ptr[t*gnb0 + col*gnb1], 50.0f)) : 0.0f;
+            decay = __shfl_sync(0xFFFFFFFF, decay, 0, WARP_SIZE);
+            state_local[i] *= decay;
+            sum1 += state_local[i] * sK[col];
+            sum2 += state_local[i] * sQ[col];
+        }
+        all_sum1[col_idx_0*WARP_SIZE_S + row] = sum1;
+        all_sum2[col_idx_0*WARP_SIZE_S + row] = sum2;
+        __syncthreads();
+
+        sum1 = sum2 = 0;
+#pragma unroll
+        for (int32_t i = 0; i < block_size/WARP_SIZE; ++i) {
+            sum1 += all_sum1[i*WARP_SIZE_S + row];
+            sum2 += all_sum2[i*WARP_SIZE_S + row];
+        }
+
+        float sv_new = beta_val * (v_ptr[t * vnb1 + row_out] - sum1);
+        if (col_idx_0 == 0) {
+            out_base[t * out_token_stride + row_out] = sum2 + sv_new * attn_score;
+        }
+
+        for (int32_t i = 0; i < HEAD_DIM/num_warps; ++i) {
+            int32_t col = num_warps*i + col_idx_0;
+            float new_state_val = state_local[i] + sv_new * sK[col];
+            new_state_val = fminf(fmaxf(new_state_val, -1e6f), 1e6f);
+            state_local[i] = new_state_val;
+        }
+
+        // Save per-step state if requested
+        if (saved_states && t < n_tokens - 1) {
+            float * state_step_dst = saved_states + batch_idx * state_batch_stride + state_head_offset + t * state_step_stride;
+            for (int32_t i = 0; i < HEAD_DIM/num_warps; ++i) {
+                int32_t col = num_warps*i + col_idx_0;
+                state_step_dst[col*HEAD_DIM + row_out] = state_local[i];
+            }
+        }
+
+        // Barrier required: (a) sK reads in the state update above must complete
+        // before next iteration overwrites sK at the top of the loop, and (b) this
+        // single barrier also orders all_sum1/all_sum2 reads above vs. the next
+        // iteration's writes — subsuming the prior barriers after the cross-warp
+        // reduction and after the loop exit.
+        __syncthreads();
+    }
+    // Copy the final state to its destination
+    for (int32_t i = 0; i < HEAD_DIM/num_warps; ++i) {
+        int32_t col = num_warps*i + col_idx_0;
+        state_dst[col*HEAD_DIM + row_out] = state_local[i];
+    }
+}
+
+static void kda_f32_cuda(
+    const float * q,
+    const float * k,
+    const float * v,
+    const float * g,
+    const float * beta,
+    const float * state_in,
+    float * dst,
+    float * state_out,
+    float * saved_states,
+    const int32_t head_dim,
+    const int32_t n_tokens,
+    const int32_t n_heads,
+    const int32_t gqa_ratio,
+    const int32_t repeat_type,
+    const int32_t n_seqs,
+    size_t vnb1, size_t vnb2, size_t vnb3,
+    size_t gnb0, size_t gnb1, size_t gnb2, size_t gnb3,
+    size_t bnb1, size_t bnb2, size_t bnb3,
+    cudaStream_t stream) {
+    if (head_dim != 64 && head_dim != 128) {
+        GGML_ABORT("Unsupported KDA head size");
+    }
+
+    const int32_t num_blocks = n_seqs * n_heads * (head_dim/WARP_SIZE);
+    const size_t smem_size = 2 * head_dim * sizeof(float);
+
+    if (n_tokens <= 8) {
+        constexpr int32_t threads_per_block = 256;
+        if (head_dim == 64) {
+            kda_recurrent_f32<64, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
+                    q, k, v, g, beta, state_in, dst, state_out, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
+                    vnb1, vnb2, vnb3, gnb0, gnb1, gnb2, gnb3, bnb1, bnb2, bnb3);
+        } else {
+            kda_recurrent_f32<128, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
+                    q, k, v, g, beta, state_in, dst, state_out, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
+                    vnb1, vnb2, vnb3, gnb0, gnb1, gnb2, gnb3, bnb1, bnb2, bnb3);
+        }
+    } else {
+        constexpr int32_t threads_per_block = 128;
+        if (head_dim == 64) {
+            kda_recurrent_f32<64, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
+                    q, k, v, g, beta, state_in, dst, state_out, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
+                    vnb1, vnb2, vnb3, gnb0, gnb1, gnb2, gnb3, bnb1, bnb2, bnb3);
+        } else {
+            kda_recurrent_f32<128, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
+                    q, k, v, g, beta, state_in, dst, state_out, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
+                    vnb1, vnb2, vnb3, gnb0, gnb1, gnb2, gnb3, bnb1, bnb2, bnb3);
+        }
+    }
+
+    CUDA_CHECK(cudaGetLastError());
+}
+
+void ggml_cuda_op_kda(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * src2 = dst->src[2];
+    const ggml_tensor * src3 = dst->src[3];
+    const ggml_tensor * src4 = dst->src[4];
+    const ggml_tensor * src5 = dst->src[5];
+    const ggml_tensor * src6 = dst->src[6];
+    const ggml_tensor * src7 = dst->src[7];
+
+    const int32_t head_dim = (int32_t) src0->ne[0];
+    const int32_t n_tokens = (int32_t) src0->ne[1];
+    const int32_t n_heads = (int32_t) src2->ne[2];
+    const int32_t n_heads_kq = (int32_t) src0->ne[2];
+    const int32_t n_seqs = (int32_t) src0->ne[3];
+    const int32_t gqa_ratio = n_heads / n_heads_kq;
+    const int32_t repeat_type = dst->op_params[0];
+    const size_t output_size = (size_t) head_dim * n_tokens * n_heads * n_seqs;
+
+    kda_f32_cuda(
+        (const float *)src0->data,
+        (const float *)src1->data,
+        (const float *)src2->data,
+        (const float *)src3->data,
+        (const float *)src4->data,
+        (const float *)src5->data,
+        (float *)dst->data,
+        src7 ? (float *)src7->data : (float *)dst->data + output_size,
+        src6 ? (float *)src6->data : nullptr,
+        head_dim, n_tokens, n_heads, gqa_ratio, repeat_type, n_seqs,
+        src2->nb[1]/sizeof(float), src2->nb[2]/sizeof(float), src2->nb[3]/sizeof(float),
+        src3->nb[0]/sizeof(float), src3->nb[1]/sizeof(float), src3->nb[2]/sizeof(float), src3->nb[3]/sizeof(float),
+        src4->nb[1]/sizeof(float), src4->nb[2]/sizeof(float), src4->nb[3]/sizeof(float),
+        ctx.stream());
+}

ggml/src/ggml-cuda/kda.cuh

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

ggml/src/ggml.c

@@ -10190,7 +10190,7 @@ struct ggml_tensor * ggml_delta_net(
 
     GGML_ASSERT(k->ne[0] == S_k && k->ne[1] == n_tokens && k->ne[2] == H_k && k->ne[3] == n_seqs);
     GGML_ASSERT(v->ne[1] == n_tokens && v->ne[3] == n_seqs);
-    GGML_ASSERT(g->ne[0] == n_tokens && g->ne[1] == 1 && g->ne[2] == H_v && g->ne[3] == n_seqs);
+    GGML_ASSERT(g->ne[0] == n_tokens && (g->ne[1] == 1 || g->ne[1] == S_v) && g->ne[2] == H_v && g->ne[3] == n_seqs);
     GGML_ASSERT(beta->ne[0] == 1 && beta->ne[1] == n_tokens && beta->ne[2] == H_v && beta->ne[3] == n_seqs);
     GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v * H_v && state->ne[2] == 1 && state->ne[3] == n_seqs);
     //GGML_ASSERT(H_k == H_v);
@@ -23863,6 +23863,148 @@ static void ggml_compute_forward_solve_tri(const struct ggml_compute_params * pa
 
 // ggml_compute_forward_delta_net
 
+static void ggml_compute_forward_kda_f32(int32_t ith, int32_t nth, struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0];
+    const struct ggml_tensor * src1 = dst->src[1];
+    const struct ggml_tensor * src2 = dst->src[2];
+    const struct ggml_tensor * src3 = dst->src[3];
+    const struct ggml_tensor * src4 = dst->src[4];
+    const struct ggml_tensor * src5 = dst->src[5];
+    const struct ggml_tensor * src6 = dst->src[6];
+    const struct ggml_tensor * src7 = dst->src[7];
+
+    const int32_t head_dim = (int32_t) src0->ne[0];
+    const int32_t n_tokens = (int32_t) src0->ne[1];
+    const int32_t n_heads = (int32_t) src2->ne[2];
+    const int32_t n_seqs = (int32_t) src0->ne[3];
+    const int32_t gqa_ratio = n_heads/(int32_t) src0->ne[2];
+
+    const size_t output_size = (size_t) head_dim * n_tokens * n_heads * n_seqs;
+
+    const float * q_data = (const float *) src0->data;
+    const float * k_data = (const float *) src1->data;
+    const float * v_data = (const float *) src2->data;
+    const float * g_data = (const float *) src3->data;
+    const float * beta_data = (const float *) src4->data;
+    const float * state_in = (const float *) src5->data;
+    float * out_data = (float *) dst->data;
+    float * saved_steps = src6 ? (float *) src6->data : NULL;
+
+    const size_t vnb1 = src2->nb[1]/sizeof(float);
+    const size_t vnb2 = src2->nb[2]/sizeof(float);
+    const size_t vnb3 = src2->nb[3]/sizeof(float);
+    const size_t gnb0 = src3->nb[0]/sizeof(float);
+    const size_t gnb1 = src3->nb[1]/sizeof(float);
+    const size_t gnb2 = src3->nb[2]/sizeof(float);
+    const size_t gnb3 = src3->nb[3]/sizeof(float);
+    const size_t bnb1 = src4->nb[1]/sizeof(float);
+    const size_t bnb2 = src4->nb[2]/sizeof(float);
+    const size_t bnb3 = src4->nb[3]/sizeof(float);
+
+    const int32_t repeat_type = dst->op_params[0];
+    const int32_t state_step_stride = head_dim * head_dim * n_heads * n_seqs;
+    float * state_working = src7 ? (float *) src7->data : out_data + output_size;
+
+#if GGML_USE_IQK_MULMAT
+    if (iqk_kda(head_dim, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
+                src2->nb[1], src2->nb[2], src2->nb[3], src3->nb, src4->nb,
+                q_data, k_data, v_data, g_data, beta_data, state_in,
+                out_data, state_working, saved_steps, state_step_stride, ith, nth)) {
+        return;
+    }
+#endif
+
+    const int32_t total_heads = n_heads * n_seqs;
+    const int32_t heads_per_thread = (total_heads + nth - 1) / nth;
+    const int32_t h_start = ith * heads_per_thread;
+    const int32_t h_end = (h_start + heads_per_thread < total_heads) ? h_start + heads_per_thread : total_heads;
+
+    const float scale = 1.0f / sqrtf((float) head_dim);
+
+    float * v_new_buf = (float *) malloc(4 * head_dim * sizeof(float));
+    GGML_ASSERT(v_new_buf);
+    float * kd = v_new_buf + head_dim;
+    float * decay = kd + head_dim;
+    float * qd = decay + head_dim;
+
+    for (int32_t h_idx = h_start; h_idx < h_end; ++h_idx) {
+        const int32_t batch_idx = h_idx / n_heads;
+        const int32_t head_idx = h_idx % n_heads;
+        const int32_t head_idx_kq = repeat_type == 0 ? head_idx / gqa_ratio : head_idx % (n_heads/gqa_ratio);
+
+        const int32_t qkv_head_offset_kq = batch_idx * (head_dim * n_tokens * n_heads/gqa_ratio) + head_idx_kq * (head_dim * n_tokens);
+        const int32_t qkv_token_stride = head_dim;
+        const int32_t state_head_offset = batch_idx * (head_dim * head_dim * n_heads) + head_idx * (head_dim * head_dim);
+        const int32_t out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim;
+        const int32_t out_token_stride = head_dim * n_heads;
+
+        const float * v_head = v_data + batch_idx * vnb3 + head_idx * vnb2;
+        const float * g_head = g_data + batch_idx * gnb3 + head_idx * gnb2;
+        const float * beta_head = beta_data + batch_idx * bnb3 + head_idx * bnb2;
+
+        float * state = state_working + state_head_offset;
+        for (int32_t i = 0; i < head_dim * head_dim; ++i) {
+            state[i] = state_in[state_head_offset + i];
+        }
+
+        const int32_t state_head_size = head_dim * head_dim;
+
+        for (int32_t t = 0; t < n_tokens; ++t) {
+            const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
+            const float * k_t = k_data + qkv_head_offset_kq + t * qkv_token_stride;
+            const float * v_t = v_head + t * vnb1;
+            const float * g_t = g_head + t * gnb0;
+
+            const float beta_raw = beta_head[t * bnb1];
+            const float beta_val = 1.0f / (1.0f + expf(-beta_raw));
+
+            float attn_score = 0.0f;
+            for (int32_t i = 0; i < head_dim; ++i) {
+                attn_score += k_t[i] * (q_t[i] * scale);
+            }
+
+            float * out_t = out_data + out_head_offset + t * out_token_stride;
+
+            for (int32_t col = 0; col < head_dim; ++col) {
+                decay[col] = expf(fminf(g_t[col * gnb1], 50.0f));
+                kd[col] = k_t[col] * decay[col];
+                qd[col] = q_t[col] * decay[col];
+            }
+
+            for (int32_t row = 0; row < head_dim; ++row) {
+                float v_prime = 0.0f;
+                float out_val = 0.0f;
+
+                for (int32_t col = 0; col < head_dim; ++col) {
+                    const float s = state[row + col * head_dim];
+                    v_prime += s * kd[col];
+                    out_val += s * qd[col];
+                }
+
+                const float v_new = v_t[row] * beta_val - v_prime * beta_val;
+                v_new_buf[row] = v_new;
+                out_t[row] = out_val * scale + v_new * attn_score;
+            }
+
+            for (int32_t col = 0; col < head_dim; ++col) {
+                const float k_col = k_t[col];
+                for (int32_t row = 0; row < head_dim; ++row) {
+                    float s = state[row + col * head_dim];
+                    s = decay[col] * s + v_new_buf[row] * k_col;
+                    state[row + col * head_dim] = fminf(fmaxf(s, -1e6f), 1e6f);
+                }
+            }
+
+            if (saved_steps && t + 1 < n_tokens) {
+                float * next_state = saved_steps + state_head_offset + (size_t) t * state_step_stride;
+                memcpy(next_state, state, state_head_size * sizeof(float));
+            }
+        }
+    }
+
+    free(v_new_buf);
+}
+
 static void ggml_compute_forward_delta_net_f32(
         const struct ggml_compute_params * params,
         struct ggml_tensor * dst) {
@@ -23875,6 +24017,11 @@ static void ggml_compute_forward_delta_net_f32(
     const struct ggml_tensor * src6 = dst->src[6];
     const struct ggml_tensor * src7 = dst->src[7];
 
+    if (src3->ne[1] != 1) {
+        ggml_compute_forward_kda_f32(params->ith, params->nth, dst);
+        return;
+    }
+
     const int64_t head_dim = src0->ne[0];
     const int64_t n_tokens = src0->ne[1];
     const int64_t n_heads  = src2->ne[2];

ggml/src/iqk/iqk_cpu_ops.h

@@ -7,6 +7,7 @@
 #pragma once
 #include <stdint.h>
 #include <stdbool.h>
+#include <stddef.h>
 #include "iqk_config.h"
 #ifdef __cplusplus
 extern "C" {
@@ -45,8 +46,14 @@ void iqk_blend(struct ggml_tensor * dst, int ith, int nth);
 
 void iqk_mask_topk(struct ggml_tensor * dst, int ith, int nth);
 
+IQK_API bool iqk_kda(int32_t head_dim, int32_t n_heads, int32_t gqa_ratio, int32_t repeat_type,
+        int32_t n_tokens, int32_t n_seqs, size_t vnb1, size_t vnb2, size_t vnb3,
+        const size_t gnb[4], const size_t bnb[4],
+        const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
+        const float * state_in, float * out_data, float * state_out, float * saved_steps,
+        int32_t state_step_stride, int32_t ith, int32_t nth);
+
 #ifdef __cplusplus
 }
 #endif
 
-

ggml/src/iqk/iqk_kda.cpp

@@ -0,0 +1,363 @@
+#include "iqk_config.h"
+#include "iqk_cpu_ops.h"
+
+#if defined IQK_IMPLEMENT
+
+#include "iqk_common.h"
+
+#include <cmath>
+#include <cstring>
+
+namespace {
+
+struct kda_layout {
+    size_t v[3];
+    size_t g[4];
+    size_t beta[4];
+};
+
+#ifdef __ARM_NEON
+template <int head_dim>
+void iqk_kda_neon_impl(int32_t n_heads, int32_t gqa_ratio, int32_t repeat_type, int32_t n_tokens, int32_t n_seqs,
+        const kda_layout & layout,
+        const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
+        const float * state_in, float * out_data, float * state_out, float * saved_steps,
+        int32_t state_step_stride, int32_t ith, int32_t nth) {
+    const int32_t total_heads = n_heads * n_seqs;
+    const int32_t heads_per_thread = (total_heads + nth - 1) / nth;
+    const int32_t h_start = ith * heads_per_thread;
+    const int32_t h_end = (h_start + heads_per_thread < total_heads) ? h_start + heads_per_thread : total_heads;
+
+    static_assert(head_dim % 4 == 0);
+
+    const float scale = 1.0f / sqrtf((float) head_dim);
+
+    float v_new_buf[head_dim];
+    float v_prime[head_dim], out_val[head_dim];
+    float decay_buf[head_dim], kd[head_dim], qd[head_dim];
+
+    float32x4x4_t vs4[4];
+
+    for (int32_t h_idx = h_start; h_idx < h_end; ++h_idx) {
+        const int32_t batch_idx = h_idx / n_heads;
+        const int32_t head_idx  = h_idx % n_heads;
+        const int32_t head_idx_kq = repeat_type == 0 ? head_idx / gqa_ratio : head_idx % (n_heads/gqa_ratio);
+
+        const int32_t state_head_offset = batch_idx * (head_dim * head_dim * n_heads) + head_idx * (head_dim * head_dim);
+        const int32_t out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim;
+        const int32_t out_token_stride = head_dim * n_heads;
+
+        const int32_t qkv_head_offset_kq = batch_idx * (head_dim * n_tokens * n_heads/gqa_ratio) + head_idx_kq * (head_dim * n_tokens);
+        const int32_t qkv_token_stride = head_dim;
+        const float * v_head = v_data + batch_idx * layout.v[2] + head_idx * layout.v[1];
+        const float * g_head = g_data + batch_idx * layout.g[3] + head_idx * layout.g[2];
+        const float * beta_head = beta_data + batch_idx * layout.beta[3] + head_idx * layout.beta[2];
+
+        float * state = state_out + state_head_offset;
+        for (int32_t i = 0; i < head_dim * head_dim; ++i) {
+            state[i] = state_in[state_head_offset + i];
+        }
+
+        for (int32_t t = 0; t < n_tokens; ++t) {
+            const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
+            const float * k_t = k_data + qkv_head_offset_kq + t * qkv_token_stride;
+            const float * v_t = v_head + t * layout.v[0];
+            const float * g_t = g_head + t * layout.g[0];
+
+            const float beta_raw = beta_head[t * layout.beta[1]];
+
+            float kq_sum = 0.0f;
+            auto vqksum = vdupq_n_f32(0.0f);
+            for (int32_t i = 0; i < head_dim; i += 4) {
+                auto vq = vld1q_f32(q_t + i);
+                auto vk = vld1q_f32(k_t + i);
+                vqksum = vfmaq_f32(vqksum, vq, vk);
+            }
+            kq_sum = vaddvq_f32(vqksum);
+
+            const float beta_val = 1.0f / (1.0f + expf(-beta_raw));
+            for (int32_t col = 0; col < head_dim; ++col) {
+                decay_buf[col] = expf(fminf(g_t[col * layout.g[1]], 50.0f));
+                kd[col] = k_t[col] * decay_buf[col];
+                qd[col] = q_t[col] * decay_buf[col];
+            }
+
+            const float attn_score = kq_sum * scale;
+
+            float * out_t = out_data + out_head_offset + t * out_token_stride;
+
+            std::memset(v_prime, 0, head_dim*sizeof(float));
+            std::memset(out_val, 0, head_dim*sizeof(float));
+            for (int32_t col = 0; col < head_dim; ++col) {
+                const float k_col = kd[col];
+                const float q_col = qd[col];
+                for (int32_t row = 0; row < head_dim; ++row) {
+                    const float s = state[row + col * head_dim];
+                    v_prime[row] += s * k_col;
+                    out_val[row] += s * q_col;
+                }
+            }
+            for (int32_t row = 0; row < head_dim; ++row) {
+                const float v_new = v_t[row] * beta_val - v_prime[row] * beta_val;
+                v_new_buf[row] = v_new;
+                out_t[row] = out_val[row] * scale + v_new * attn_score;
+            }
+
+            auto vmin = vdupq_n_f32(-1e6f);
+            auto vmax = vdupq_n_f32( 1e6f);
+            for (int32_t col = 0; col < head_dim; col += 4) {
+                auto vk = vld1q_f32(k_t + col);
+                for (int32_t row = 0; row < head_dim; row += 16) {
+                    for (int32_t k = 0; k < 4; ++k) {
+                        vs4[k] = vld1q_f32_x4(state + (col + k)*head_dim + row);
+                        auto vd = vdupq_n_f32(decay_buf[col + k]);
+                        for (int32_t j = 0; j < 4; ++j) {
+                            vs4[k].val[j] = vmulq_f32(vs4[k].val[j], vd);
+                        }
+                    }
+                    auto vn = vld1q_f32_x4(v_new_buf + row);
+                    for (int32_t j = 0; j < 4; ++j) {
+                        vs4[0].val[j] = vfmaq_laneq_f32(vs4[0].val[j], vn.val[j], vk, 0);
+                        vs4[1].val[j] = vfmaq_laneq_f32(vs4[1].val[j], vn.val[j], vk, 1);
+                        vs4[2].val[j] = vfmaq_laneq_f32(vs4[2].val[j], vn.val[j], vk, 2);
+                        vs4[3].val[j] = vfmaq_laneq_f32(vs4[3].val[j], vn.val[j], vk, 3);
+                    }
+                    for (int32_t k = 0; k < 4; ++k) {
+                        for (int32_t j = 0; j < 4; ++j) {
+                            vs4[k].val[j] = vmaxq_f32(vminq_f32(vs4[k].val[j], vmax), vmin);
+                        }
+                        vst1q_f32_x4(state + (col + k)*head_dim + row, vs4[k]);
+                    }
+                }
+            }
+
+            if (saved_steps && t + 1 < n_tokens) {
+                float * this_state = saved_steps + state_head_offset + t * state_step_stride;
+                std::memcpy(this_state, state, head_dim * head_dim * sizeof(float));
+            }
+        }
+    }
+}
+#endif
+
+template <int head_dim>
+void iqk_kda_impl(int32_t n_heads, int32_t gqa_ratio, int32_t repeat_type, int32_t n_tokens, int32_t n_seqs,
+        const kda_layout & layout,
+        const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
+        const float * state_in, float * out_data, float * state_out, float * saved_steps,
+        int32_t state_step_stride, int32_t ith, int32_t nth) {
+#ifdef __ARM_NEON
+    iqk_kda_neon_impl<head_dim>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, layout,
+            q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out,
+            saved_steps, state_step_stride, ith, nth);
+    return;
+#endif
+    const int32_t total_heads = n_heads * n_seqs;
+    const int32_t heads_per_thread = (total_heads + nth - 1) / nth;
+    const int32_t h_start = ith * heads_per_thread;
+    const int32_t h_end = (h_start + heads_per_thread < total_heads) ? h_start + heads_per_thread : total_heads;
+
+#ifdef __AVX2__
+    static_assert(head_dim % 8 == 0);
+#endif
+
+    const float scale = 1.0f / sqrtf((float) head_dim);
+
+#ifdef __AVX512F__
+    __m512 v_prime[head_dim/16], out_val[head_dim/16];
+#else
+    float v_new_buf[head_dim];
+    float v_prime[head_dim], out_val[head_dim];
+#endif
+    float decay_buf[head_dim];
+
+    for (int32_t h_idx = h_start; h_idx < h_end; ++h_idx) {
+        const int32_t batch_idx = h_idx / n_heads;
+        const int32_t head_idx = h_idx % n_heads;
+        const int32_t head_idx_kq = repeat_type == 0 ? head_idx / gqa_ratio : head_idx % (n_heads/gqa_ratio);
+
+        const int32_t state_head_offset = batch_idx * (head_dim * head_dim * n_heads) + head_idx * (head_dim * head_dim);
+        const int32_t out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim;
+        const int32_t out_token_stride = head_dim * n_heads;
+
+        const int32_t qkv_head_offset_kq = batch_idx * (head_dim * n_tokens * n_heads/gqa_ratio) + head_idx_kq * (head_dim * n_tokens);
+        const int32_t qkv_token_stride = head_dim;
+        const float * v_head = v_data + batch_idx * layout.v[2] + head_idx * layout.v[1];
+        const float * g_head = g_data + batch_idx * layout.g[3] + head_idx * layout.g[2];
+        const float * beta_head = beta_data + batch_idx * layout.beta[3] + head_idx * layout.beta[2];
+
+        float * state = state_out + state_head_offset;
+        for (int32_t i = 0; i < head_dim * head_dim; ++i) {
+            state[i] = state_in[state_head_offset + i];
+        }
+
+        for (int32_t t = 0; t < n_tokens; ++t) {
+            const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
+            const float * k_t = k_data + qkv_head_offset_kq + t * qkv_token_stride;
+            const float * v_t = v_head + t * layout.v[0];
+            const float * g_t = g_head + t * layout.g[0];
+
+            const float beta_raw = beta_head[t * layout.beta[1]];
+
+            float kq_sum = 0.0f;
+#if defined __AVX512F__
+            auto vqksum = _mm512_setzero_ps();
+            for (int32_t i = 0; i < head_dim; i += 16) {
+                auto vq = _mm512_loadu_ps(q_t + i);
+                auto vk = _mm512_loadu_ps(k_t + i);
+                vqksum = _mm512_fmadd_ps(vk, vq, vqksum);
+            }
+            kq_sum = _mm512_reduce_add_ps(vqksum);
+#elif defined __AVX2__
+            auto vqksum = _mm256_setzero_ps();
+            for (int32_t i = 0; i < head_dim; i += 8) {
+                auto vq = _mm256_loadu_ps(q_t + i);
+                auto vk = _mm256_loadu_ps(k_t + i);
+                vqksum = _mm256_fmadd_ps(vk, vq, vqksum);
+            }
+            kq_sum = hsum_float_8(vqksum);
+#else
+            for (int32_t i = 0; i < head_dim; ++i) {
+                kq_sum += k_t[i] * q_t[i];
+            }
+#endif
+
+            const float beta_val = 1.0f / (1.0f + expf(-beta_raw));
+            for (int32_t col = 0; col < head_dim; ++col) {
+                decay_buf[col] = expf(fminf(g_t[col * layout.g[1]], 50.0f));
+            }
+
+            const float attn_score = kq_sum * scale;
+
+            float * out_t = out_data + out_head_offset + t * out_token_stride;
+
+#ifdef __AVX512F__
+            for (int32_t j = 0; j < head_dim/16; ++j) {
+                v_prime[j] = out_val[j] = _mm512_setzero_ps();
+            }
+            for (int32_t col = 0; col < head_dim; ++col) {
+                auto k_col = _mm512_set1_ps(k_t[col] * decay_buf[col]);
+                auto q_col = _mm512_set1_ps(q_t[col] * decay_buf[col]);
+                for (int32_t j = 0; j < head_dim/16; ++j) {
+                    auto s = _mm512_loadu_ps(state + col * head_dim + 16*j);
+                    v_prime[j] = _mm512_fmadd_ps(s, k_col, v_prime[j]);
+                    out_val[j] = _mm512_fmadd_ps(s, q_col, out_val[j]);
+                }
+            }
+            auto c1 = _mm512_set1_ps(beta_val);
+            auto c2 = _mm512_set1_ps(beta_val);
+            auto c3 = _mm512_set1_ps(scale);
+            auto c4 = _mm512_set1_ps(attn_score);
+            for (int32_t j = 0; j < head_dim/16; ++j) {
+                auto v = _mm512_loadu_ps(v_t + 16*j);
+                v_prime[j] = _mm512_sub_ps(_mm512_mul_ps(v, c1), _mm512_mul_ps(v_prime[j], c2));
+                auto oval = _mm512_fmadd_ps(v_prime[j], c4, _mm512_mul_ps(out_val[j], c3));
+                _mm512_storeu_ps(out_t + 16*j, oval);
+            }
+            auto vmin = _mm512_set1_ps(-1e6f);
+            auto vmax = _mm512_set1_ps( 1e6f);
+            for (int32_t col = 0; col < head_dim; ++col) {
+                auto vk = _mm512_set1_ps(k_t[col]);
+                auto vd = _mm512_set1_ps(decay_buf[col]);
+                for (int32_t j = 0; j < head_dim/16; ++j) {
+                    auto vs = _mm512_loadu_ps(state + col * head_dim + 16*j);
+                    vs = _mm512_fmadd_ps(v_prime[j], vk, _mm512_mul_ps(vs, vd));
+                    vs = _mm512_max_ps(vmin, _mm512_min_ps(vmax, vs));
+                    _mm512_storeu_ps(state + col * head_dim + 16*j, vs);
+                }
+            }
+#else
+            std::memset(v_prime, 0, head_dim*sizeof(float));
+            std::memset(out_val, 0, head_dim*sizeof(float));
+            for (int32_t col = 0; col < head_dim; ++col) {
+                const float k_col = k_t[col] * decay_buf[col];
+                const float q_col = q_t[col] * decay_buf[col];
+                for (int32_t row = 0; row < head_dim; ++row) {
+                    const float s = state[row + col * head_dim];
+                    v_prime[row] += s * k_col;
+                    out_val[row] += s * q_col;
+                }
+            }
+            for (int32_t row = 0; row < head_dim; ++row) {
+                const float v_new = v_t[row] * beta_val - v_prime[row] * beta_val;
+                v_new_buf[row] = v_new;
+                out_t[row] = out_val[row] * scale + v_new * attn_score;
+            }
+
+#ifdef __AVX2__
+            auto vmin = _mm256_set1_ps(-1e6f);
+            auto vmax = _mm256_set1_ps( 1e6f);
+            for (int32_t col = 0; col < head_dim; ++col) {
+                auto vk = _mm256_set1_ps(k_t[col]);
+                auto vd = _mm256_set1_ps(decay_buf[col]);
+                for (int32_t row = 0; row < head_dim; row += 8) {
+                    auto vs = _mm256_loadu_ps(state + col * head_dim + row);
+                    auto vn = _mm256_loadu_ps(v_new_buf + row);
+                    vs = _mm256_fmadd_ps(vn, vk, _mm256_mul_ps(vs, vd));
+                    auto mask_l = _mm256_cmp_ps(vs, vmin, _CMP_LT_OQ);
+                    auto mask_u = _mm256_cmp_ps(vs, vmax, _CMP_GT_OQ);
+                    vs = _mm256_or_ps(_mm256_and_ps(mask_l, vmin), _mm256_andnot_ps(mask_l, vs));
+                    vs = _mm256_or_ps(_mm256_and_ps(mask_u, vmax), _mm256_andnot_ps(mask_u, vs));
+                    _mm256_storeu_ps(state + col * head_dim + row, vs);
+                }
+            }
+#else
+            for (int32_t col = 0; col < head_dim; ++col) {
+                const float k_col = k_t[col];
+                for (int32_t row = 0; row < head_dim; ++row) {
+                    float s = state[row + col * head_dim];
+                    s = decay_buf[col] * s + v_new_buf[row] * k_col;
+                    state[row + col * head_dim] = fminf(fmaxf(s, -1e6f), 1e6f);
+                }
+            }
+#endif
+#endif
+
+            if (saved_steps && t + 1 < n_tokens) {
+                float * this_state = saved_steps + state_head_offset + t * state_step_stride;
+                std::memcpy(this_state, state, head_dim * head_dim * sizeof(float));
+            }
+        }
+    }
+}
+
+}
+
+extern "C" IQK_API bool iqk_kda(int32_t head_dim, int32_t n_heads, int32_t gqa_ratio, int32_t repeat_type,
+        int32_t n_tokens, int32_t n_seqs, size_t vnb1, size_t vnb2, size_t vnb3,
+        const size_t gnb[4], const size_t bnb[4],
+        const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
+        const float * state_in, float * out_data, float * state_out, float * saved_steps,
+        int32_t state_step_stride, int32_t ith, int32_t nth) {
+    if (head_dim != 64 && head_dim != 128) {
+        return false;
+    }
+
+    const kda_layout layout = {
+        { vnb1/sizeof(float), vnb2/sizeof(float), vnb3/sizeof(float) },
+        { gnb[0]/sizeof(float), gnb[1]/sizeof(float), gnb[2]/sizeof(float), gnb[3]/sizeof(float) },
+        { bnb[0]/sizeof(float), bnb[1]/sizeof(float), bnb[2]/sizeof(float), bnb[3]/sizeof(float) },
+    };
+
+    if (head_dim == 64) {
+        iqk_kda_impl<64>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, layout,
+                q_data, k_data, v_data, g_data, beta_data, state_in,
+                out_data, state_out, saved_steps, state_step_stride, ith, nth);
+    } else {
+        iqk_kda_impl<128>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, layout,
+                q_data, k_data, v_data, g_data, beta_data, state_in,
+                out_data, state_out, saved_steps, state_step_stride, ith, nth);
+    }
+    return true;
+}
+
+#else
+
+extern "C" IQK_API bool iqk_kda(int32_t, int32_t, int32_t, int32_t, int32_t, int32_t,
+        size_t, size_t, size_t, const size_t[4], const size_t[4],
+        const float *, const float *, const float *, const float *, const float *,
+        const float *, float *, float *, float *, int32_t, int32_t, int32_t) {
+    return false;
+}
+
+#endif

src/CMakeLists.txt

@@ -63,6 +63,8 @@ add_library(llama
             llama-hparams.cpp
             llama-delta-net.h
             llama-delta-net.cpp
+            llama-kda.h
+            llama-kda.cpp
             unicode.h
             unicode.cpp
             unicode-data.cpp
@@ -128,6 +130,7 @@ add_library(llama
             graphs/build_mimo2.cpp
             graphs/build_openai.cpp
             graphs/build_bailingmoe2.cpp
+            graphs/build_bailingmoe3.cpp
             graphs/build_minimaxm2.cpp
             graphs/build_minimaxm3.cpp
             graphs/build_smollm3.cpp

src/graphs/build_bailingmoe3.cpp

@@ -0,0 +1,99 @@
+#include "../llama-build-context.h"
+#include "../llama-context.h"
+#include "../llama-delta-net.h"
+#include "../llama-model.h"
+
+ggml_cgraph * llm_build_context::build_bailingmoe3() {
+    const bool tp_mode = model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN;
+#ifdef GGML_USE_VULKAN
+    const bool use_f32_attn_precision = true;
+#else
+    const bool use_f32_attn_precision = lctx.cparams.graph_attn_precision == GGML_TYPE_F32;
+#endif
+
+    ggml_cgraph * gf = new_graph_custom();
+    delta_net delta(lctx, batch);
+
+    auto inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
+    auto inp_pos = build_inp_pos();
+    auto inp_out_ids = build_inp_out_ids();
+    auto KQ_mask = build_inp_KQ_mask();
+
+    lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens);
+    cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1);
+    ggml_set_input(lctx.inp_s_seq_qnext);
+
+    const float mscale = attn_factor * (1.0f + hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
+    const float kq_scale = mscale * mscale / sqrtf(float(hparams.n_embd_head_k(0)));
+    const float attn_factor_scaled = 1.0f / (1.0f + 0.1f * logf(1.0f / freq_scale));
+    const bool pp_opt = n_tokens >= 128 && lctx.cparams.mla_attn > 1;
+    auto rope_cache = cparams.rope_cache && (rope_type == LLAMA_ROPE_TYPE_NEOX || rope_type == LLAMA_ROPE_TYPE_NORM)
+        ? ggml_rope_cache(ctx0, inp_pos, nullptr, n_rot, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow)
+        : nullptr;
+
+    const int n_transformer_layers = hparams.n_layer_kv_from_start;
+    ggml_tensor * cur = nullptr;
+    for (int il = 0; il < n_transformer_layers; ++il) {
+        const bool last_layer = il == n_transformer_layers - 1;
+
+        if (hparams.is_recurrent(il)) {
+            cur = delta.build_layer_attn_kda(ctx0, gf, inpL, last_layer ? inp_out_ids : nullptr, il, cb);
+        } else {
+            auto inpSA = inpL;
+            const bool is_tp_layer = tp_mode && model.layers[il].wo->extra;
+            const bool direct_q = hparams.n_lora_q == 0;
+            if (is_tp_layer) {
+                cur = build_deepseek2_tp_attention(gf, il, inpL, KQ_mask, inp_pos, rope_cache,
+                        kq_scale, attn_factor_scaled, use_f32_attn_precision, direct_q, pp_opt);
+            } else {
+                cur = build_deepseek2_layer_attention(gf, il, inpL, KQ_mask, inp_pos, rope_cache,
+                        kq_scale, attn_factor_scaled, use_f32_attn_precision, direct_q, pp_opt);
+            }
+            if (last_layer) {
+                cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+                if (!is_tp_layer) {
+                    inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+                }
+            }
+            if (!is_tp_layer) {
+                cur = ggml_add(ctx0, cur, inpSA);
+            }
+            cb(cur, "ffn_inp", il);
+        }
+
+        if ((uint32_t) il < hparams.n_layer_dense_lead) {
+            cur = llm_build_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur,
+                    model.layers[il].ffn_up,   nullptr, nullptr,
+                    model.layers[il].ffn_gate, nullptr, nullptr,
+                    model.layers[il].ffn_down, nullptr, nullptr,
+                    nullptr,
+                    LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, true);
+        } else {
+            cur = llm_build_std_moe_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur,
+                    model.layers[il].ffn_gate_inp, nullptr,
+                    model.layers[il].ffn_up_exps, nullptr,
+                    model.layers[il].ffn_gate_exps, nullptr,
+                    model.layers[il].ffn_down_exps, nullptr,
+                    model.layers[il].ffn_exp_probs_b,
+                    model.layers[il].ffn_up_shexp, nullptr,
+                    model.layers[il].ffn_gate_shexp, nullptr,
+                    model.layers[il].ffn_down_shexp, nullptr,
+                    n_expert, n_expert_used,
+                    LLM_FFN_SILU, hparams.expert_weights_norm,
+                    true, hparams.expert_weights_scale,
+                    (llm_expert_gating_func_type) hparams.expert_gating_func,
+                    LLM_FFN_SILU, cb, il, gf, true, model.layers[il].ffn_up_gate_exps);
+        }
+        cb(cur, "ffn_out", il);
+
+        cur = lctx.cvec.apply_to(ctx0, cur, il);
+        cb(cur, "l_out", il);
+        inpL = cur;
+    }
+
+    cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb);
+    cb(cur, "result_output", -1);
+    ggml_build_forward_expand(gf, cur);
+    return gf;
+}

src/graphs/build_deepseek2.cpp

@@ -302,6 +302,13 @@ ggml_tensor * llm_build_context::build_deepseek2_tp_attention(
             kqv_2d = ggml_reshape_2d(ctx0, kqv, n_embd_head_v * n_head_local, n_tokens);
         }
 
+        if (model.layers[il].wqkv_gate) {
+            auto gate_split = (const ggml_split_tensor_t *) model.layers[il].wqkv_gate->extra;
+            GGML_ASSERT(gate_split && gate_split->splits[id]);
+            kqv_2d = build_mla_output_gate(lctx, ctx0, kqv_2d, cur, gate_split->splits[id],
+                    n_embd_head_v, n_head_local, il_id, cb);
+        }
+
         ggml_tensor * partial = llm_build_lora_mm(lctx, ctx0, wo_split->splits[id], kqv_2d);
 
         // Fold residual into the first non-skipped rank so the reduce result includes it.
@@ -728,6 +735,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
     // norm
     cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il);
     cb(cur, "attn_norm", il);
+    auto input_normed = cur;
 
     // DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). Built below from the q_lora latent
     // and used to construct a sparse causal mask. Defaults to the dense KQ_mask.
@@ -1160,6 +1168,10 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
 
             }
 
+            if (model.layers[il].wqkv_gate) {
+                cur = build_mla_output_gate(lctx, ctx0, cur, input_normed, model.layers[il].wqkv_gate,
+                        n_embd_head_v, n_head, il, cb);
+            }
             ggml_build_forward_expand(gf, cur);
 
             cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wo, cur);
@@ -1200,9 +1212,19 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
             struct ggml_tensor * k_states = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_rope, q_rope), 0);
             cb(k_states, "k_states", il);
 
-            cur = llm_build_kv(ctx0, lctx, kv_self, gf,
-                    model.layers[il].wo, NULL,
-                    k_states, v_states, q_states, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il);
+            if (model.layers[il].wqkv_gate) {
+                cur = llm_build_kv(ctx0, lctx, kv_self, gf,
+                        nullptr, nullptr,
+                        k_states, v_states, q_states, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il);
+                cur = build_mla_output_gate(lctx, ctx0, cur, input_normed, model.layers[il].wqkv_gate,
+                        n_embd_head_v, n_head, il, cb);
+                cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wo, cur);
+                cb(cur, "kqv_out", il);
+            } else {
+                cur = llm_build_kv(ctx0, lctx, kv_self, gf,
+                        model.layers[il].wo, nullptr,
+                        k_states, v_states, q_states, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il);
+            }
 
         }
     }

src/llama-arch.cpp

@@ -73,6 +73,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_HUNYUAN_MOE,     "hunyuan-moe"  },
     { LLM_ARCH_OPENAI_MOE,      "gpt-oss"      },
     { LLM_ARCH_BAILINGMOE2,     "bailingmoe2"  },
+    { LLM_ARCH_BAILINGMOE3,     "bailingmoe3"  },
     { LLM_ARCH_MINIMAX_M2,      "minimax-m2"   },
     { LLM_ARCH_MINIMAX_M3,      "minimax-m3"   },
     { LLM_ARCH_SMOLLM3,         "smollm3"      },
@@ -241,6 +242,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_SSM_STATE_SIZE,                "%s.ssm.state_size"     },
     { LLM_KV_SSM_TIME_STEP_RANK,            "%s.ssm.time_step_rank" },
     { LLM_KV_SSM_GROUP_COUNT,               "%s.ssm.group_count"    },
+    { LLM_KV_KDA_HEAD_DIM,                  "%s.kda.head_dim"       },
+    { LLM_KV_KDA_SAFE_GATE,                 "%s.kda.safe_gate"      },
+    { LLM_KV_KDA_GATE_LOWER_BOUND,          "%s.kda.gate_lower_bound" },
 
     { LLM_KV_TOKENIZER_MODEL,                "tokenizer.ggml.model"                    },
     { LLM_KV_TOKENIZER_PRE,                  "tokenizer.ggml.pre"                      },
@@ -316,6 +320,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
     case LLM_ARCH_QWEN3NEXT:
     case LLM_ARCH_QWEN35MOE:
     case LLM_ARCH_QWEN35:
+    case LLM_ARCH_BAILINGMOE3:
         return true;
     default:
         return false;

src/llama-arch.h

@@ -71,6 +71,7 @@ enum llm_arch {
     LLM_ARCH_HUNYUAN_MOE,
     LLM_ARCH_OPENAI_MOE,
     LLM_ARCH_BAILINGMOE2,
+    LLM_ARCH_BAILINGMOE3,
     LLM_ARCH_MINIMAX_M2,
     LLM_ARCH_MINIMAX_M3,
     LLM_ARCH_SMOLLM3,
@@ -225,6 +226,9 @@ enum llm_kv {
     LLM_KV_SSM_STATE_SIZE,
     LLM_KV_SSM_TIME_STEP_RANK,
     LLM_KV_SSM_GROUP_COUNT,
+    LLM_KV_KDA_HEAD_DIM,
+    LLM_KV_KDA_SAFE_GATE,
+    LLM_KV_KDA_GATE_LOWER_BOUND,
 
     LLM_KV_TOKENIZER_MODEL,
     LLM_KV_TOKENIZER_PRE,
@@ -326,6 +330,9 @@ enum llm_tensor {
     LLM_TENSOR_LAYER_OUT_NORM,
     LLM_TENSOR_SSM_IN,
     LLM_TENSOR_SSM_CONV1D,                  // 45
+    LLM_TENSOR_SSM_CONV1D_Q,
+    LLM_TENSOR_SSM_CONV1D_K,
+    LLM_TENSOR_SSM_CONV1D_V,
     LLM_TENSOR_SSM_X,
     LLM_TENSOR_SSM_DT,
     LLM_TENSOR_SSM_A,
@@ -336,6 +343,10 @@ enum llm_tensor {
     LLM_TENSOR_SSM_BETA_ALPHA,
     LLM_TENSOR_SSM_ALPHA,
     LLM_TENSOR_SSM_BETA,                    // 50
+    LLM_TENSOR_SSM_F_A,
+    LLM_TENSOR_SSM_G_A,
+    LLM_TENSOR_SSM_F,
+    LLM_TENSOR_SSM_G,
     LLM_TENSOR_ATTN_Q_A,
     LLM_TENSOR_ATTN_Q_B,
     LLM_TENSOR_ATTN_KV_A_MQA,

src/llama-build-context.cpp

@@ -989,6 +989,27 @@ ggml_tensor * llm_build_context::llm_build_lora_mm(
     return res;
 }
 
+ggml_tensor * llm_build_context::build_mla_output_gate(
+        llama_context & lctx,
+        ggml_context * ctx0,
+        ggml_tensor * output,
+        ggml_tensor * input,
+        ggml_tensor * gate_w,
+        int64_t head_dim,
+        int64_t n_head,
+        int il,
+        const llm_build_cb & cb) {
+    auto gate = llm_build_lora_mm(lctx, ctx0, gate_w, input);
+    gate = ggml_reshape_3d(ctx0, gate, 1, n_head, output->ne[1]);
+    cb(gate, "attn_gate", il);
+
+    output = ggml_reshape_3d(ctx0, output, head_dim, n_head, output->ne[1]);
+    output = ggml_fused_mul_unary(ctx0, gate, output, GGML_UNARY_OP_SIGMOID);
+    output = ggml_reshape_2d(ctx0, output, head_dim * n_head, output->ne[2]);
+    cb(output, "attn_gated", il);
+    return output;
+}
+
 ggml_tensor * llm_build_context::llm_build_lora_mm_id(
         struct llama_context & lctx,
          struct ggml_context * ctx0,
@@ -1166,9 +1187,7 @@ ggml_tensor * llm_build_context::llm_build_ffn(
             }
             cur = ggml_fused_up_gate(ctx, split_u, split_g, cur, unary_op);
             cb(cur, "ffn_up_gate", il_cb);
-            if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_DEEPSEEK4) {
-                *(float *)(cur->op_params + 1) = lctx.model.hparams.swiglu_limits[il];
-            }
+            *(float *)(cur->op_params + 1) = lctx.model.swiglu_limit(il, lctx.model.arch == LLM_ARCH_BAILINGMOE3);
             cur = llm_build_lora_mm(lctx, ctx, split_d, cur);
             cb(cur, "ffn_down", il_cb);
             if (lctx.model.arch == LLM_ARCH_GLM4 || lctx.model.arch == LLM_ARCH_GLM4_MOE) {
@@ -1227,9 +1246,7 @@ ggml_tensor * llm_build_context::llm_build_ffn(
                         type_op == LLM_FFN_GELU ? GGML_UNARY_OP_GELU : GGML_UNARY_OP_SWIGLU_OAI;
         cur = ggml_fused_up_gate(ctx, up, gate, cur, unary_op);
         cb(cur, "ffn_up_gate", il);
-        if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_DEEPSEEK4) {
-            *(float *)(cur->op_params + 1) = lctx.model.hparams.swiglu_limits_shared[il];
-        }
+        *(float *)(cur->op_params + 1) = lctx.model.swiglu_limit(il, true);
         if (down) {
             cur = llm_build_lora_mm(lctx, ctx, down, cur);
             cb(cur, "ffn_down", il);
@@ -1310,18 +1327,16 @@ ggml_tensor * llm_build_context::llm_build_ffn(
        (type_op == LLM_FFN_SILU || type_op == LLM_FFN_RELU || (type_op == LLM_FFN_GELU && !act_scales))) {
         cur = ggml_fused_mul_unary(ctx, cur, tmp, type_op == LLM_FFN_SILU ? GGML_UNARY_OP_SILU :
                                                   type_op == LLM_FFN_RELU ? GGML_UNARY_OP_RELU : GGML_UNARY_OP_GELU);
-        if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_DEEPSEEK4) {
-            *((float *)(cur->op_params + 1)) = lctx.model.hparams.swiglu_limits_shared[il];
-        }
+        *((float *)(cur->op_params + 1)) = lctx.model.swiglu_limit(il, true);
     }
     else {
 
     switch (type_op) {
         case LLM_FFN_SILU:
             {
-                if (lctx.model.arch == LLM_ARCH_STEP35) {
+                if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_BAILINGMOE3) {
                     cur = ggml_fused_mul_unary(ctx, cur, up, GGML_UNARY_OP_SILU);
-                    *(float *)(cur->op_params + 1) = lctx.model.hparams.swiglu_limits_shared[il];
+                    *(float *)(cur->op_params + 1) = lctx.model.swiglu_limit(il, true);
                     type_gate = LLM_FFN_SEQ;
                     break;
                 }
@@ -1497,7 +1512,9 @@ llm_expert_gating_func_type   gating_op,
 
     // select experts
     if (selected_experts == nullptr) {
-        if (lctx.cparams.grouped_expert_routing && lctx.model.arch == LLM_ARCH_BAILINGMOE2 && n_tokens > 0) {
+        const bool grouped_routing = lctx.cparams.grouped_expert_routing &&
+                (lctx.model.arch == LLM_ARCH_BAILINGMOE2 || lctx.model.arch == LLM_ARCH_BAILINGMOE3);
+        if (grouped_routing && n_tokens > 0) {
             auto& hparams = lctx.model.hparams;
             selected_experts = ggml_grouped_topk(ctx, selection_probs, hparams.n_expert_groups, hparams.n_group_used, 2, n_expert_used);
         } else {
@@ -1527,7 +1544,7 @@ llm_expert_gating_func_type   gating_op,
         if (lctx.model.arch == LLM_ARCH_LAGUNA) {
             weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5f, INFINITY);
             cb(weights_sum, "ffn_moe_weights_sum_clamped", il);
-        } else if (lctx.model.arch == LLM_ARCH_BAILINGMOE2 || lctx.model.arch == LLM_ARCH_STEP35) {
+        } else if (lctx.model.arch == LLM_ARCH_BAILINGMOE2 || lctx.model.arch == LLM_ARCH_BAILINGMOE3 || lctx.model.arch == LLM_ARCH_STEP35) {
             weights_sum = ggml_scale_bias(ctx, weights_sum, 1.0, 1e-20);
             cb(weights_sum, "ffn_moe_weights_sum_biased", il);
         }
@@ -1573,9 +1590,7 @@ llm_expert_gating_func_type   gating_op,
             par = ggml_moe_up_gate(ctx, up_gate_exps, nullptr, cur, selected_experts,
                     type_op == LLM_FFN_SILU ? GGML_UNARY_OP_SILU : GGML_UNARY_OP_GELU);
         }
-        if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_DEEPSEEK4) {
-            *((float *)(par->op_params + 1)) = lctx.model.hparams.swiglu_limits[il];
-        }
+        *((float *)(par->op_params + 1)) = lctx.model.swiglu_limit(il, false);
     } else {
     GGML_ASSERT(!up_gate_exps && !up_gate_exps_b);
 
@@ -1589,9 +1604,7 @@ llm_expert_gating_func_type   gating_op,
             par = ggml_moe_up_gate(ctx, up_exps, gate_exps, cur, selected_experts,
                     type_op == LLM_FFN_SILU ? GGML_UNARY_OP_SILU : GGML_UNARY_OP_GELU);
         }
-        if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_DEEPSEEK4) {
-            *(float *)(par->op_params + 1) = lctx.model.hparams.swiglu_limits[il];
-        }
+        *(float *)(par->op_params + 1) = lctx.model.swiglu_limit(il, false);
     } else {
         ggml_tensor * up = llm_build_lora_mm_id(lctx, ctx, up_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]
         cb(up, "ffn_moe_up", il);
@@ -1617,9 +1630,7 @@ llm_expert_gating_func_type   gating_op,
 
         if (type_op == LLM_FFN_SILU || type_op == LLM_FFN_GELU) {
             par = ggml_fused_mul_unary(ctx, gate, up, type_op == LLM_FFN_SILU ? GGML_UNARY_OP_SILU : GGML_UNARY_OP_GELU);
-            if (lctx.model.arch == LLM_ARCH_STEP35 || lctx.model.arch == LLM_ARCH_DEEPSEEK4) {
-                *((float *)(par->op_params + 1)) = lctx.model.hparams.swiglu_limits[il];
-            }
+            *((float *)(par->op_params + 1)) = lctx.model.swiglu_limit(il, false);
         } else if (type_op == LLM_FFN_SWIGLU_OAI) {
             constexpr float alpha = 1.702f;
             constexpr float limit = 7.0f;
@@ -2986,6 +2997,10 @@ ggml_cgraph * llm_build_context::llama_build_graph(
             {
                 result = llm.build_bailingmoe2();
             } break;
+        case LLM_ARCH_BAILINGMOE3:
+            {
+                result = llm.build_bailingmoe3();
+            } break;
         case LLM_ARCH_MINIMAX_M2:
             {
                 result = llm.build_minimaxm2();

src/llama-build-context.h

@@ -429,6 +429,7 @@ struct llm_build_context {
     ggml_cgraph * build_openai_moe();
 
     ggml_cgraph * build_bailingmoe2();
+    ggml_cgraph * build_bailingmoe3();
 
     ggml_cgraph * build_minimaxm2();
     ggml_cgraph * build_minimaxm3();
@@ -456,6 +457,10 @@ struct llm_build_context {
     static ggml_tensor * llm_build_lora_mm(llama_context & lctx, ggml_context * ctx0,
             ggml_tensor * w, ggml_tensor * cur);
 
+    static ggml_tensor * build_mla_output_gate(llama_context & lctx, ggml_context * ctx0,
+            ggml_tensor * output, ggml_tensor * input, ggml_tensor * gate_w,
+            int64_t head_dim, int64_t n_head, int il, const llm_build_cb & cb);
+
     static ggml_tensor * llm_build_lora_mm_id(llama_context & lctx, ggml_context * ctx0,
           ggml_tensor * w, ggml_tensor * cur, ggml_tensor * ids);
 

src/llama-delta-net.cpp

@@ -95,7 +95,9 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
     GGML_ASSERT(q->ne[0] == S_k && q->ne[2] == H_k && q->ne[1] == n_tokens && q->ne[3] == n_seqs);
     GGML_ASSERT(k->ne[0] == S_k && k->ne[2] == H_k && k->ne[1] == n_tokens && k->ne[3] == n_seqs);
     GGML_ASSERT(v->ne[2] == n_tokens);
-    GGML_ASSERT(g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs);
+    const bool scalar_gate = g->ne[0] == H_v && g->ne[1] == n_tokens && g->ne[2] == n_seqs;
+    const bool channel_gate = g->ne[0] == S_v && g->ne[1] == H_v && g->ne[2] == n_tokens && g->ne[3] == n_seqs;
+    GGML_ASSERT(scalar_gate || channel_gate);
     GGML_ASSERT(beta->ne[0] == H_v && beta->ne[2] == n_tokens && beta->ne[3] == n_seqs);
     GGML_ASSERT(state->ne[0] == S_v && state->ne[1] == S_v && state->ne[2] == H_v && state->ne[3] == n_seqs);
     //GGML_ASSERT(H_k == H_v);
@@ -109,7 +111,7 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
     cb(state,"state_in", il);
 
     v = ggml_permute(ctx0, v, 0, 2, 1, 3);
-    g = ggml_permute(ctx0, g, 2, 0, 3, 1);
+    g = channel_gate ? ggml_permute(ctx0, g, 1, 2, 0, 3) : ggml_permute(ctx0, g, 2, 0, 3, 1);
     beta = ggml_permute(ctx0, beta, 2, 0, 1, 3);
 
     ggml_tensor * state_flat = ggml_reshape_4d(ctx0, state, S_v, S_v * H_v, 1, n_seqs);

src/llama-delta-net.h

@@ -24,6 +24,9 @@ struct delta_net {
     ggml_tensor * build_layer_attn_linear(ggml_context * ctx0, ggml_cgraph * gf,
             ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb) const;
 
+    ggml_tensor * build_layer_attn_kda(ggml_context * ctx0, ggml_cgraph * gf,
+            ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb) const;
+
 private:
 
     llama_context     & lctx;
@@ -59,4 +62,8 @@ private:
 
     static ggml_tensor * build_gated_output(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_norm, ggml_tensor * ssm_out,
             ggml_tensor * output, ggml_tensor * z, int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb);
+
+    ggml_tensor * build_layer_attn_kda_core(ggml_context * ctx0, ggml_cgraph * gf,
+            ggml_tensor * cur, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids,
+            uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb) const;
 };

src/llama-hparams.cpp

@@ -1401,7 +1401,7 @@ void llm_load_hparams(
                 ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS,              hparams.nextn_predict_layers, false);
 
                 // TODO: when MTP is implemented, this should probably be updated if needed
-                hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
+                hparams.n_layer_kv_from_start = static_cast<int32_t>(hparams.n_layer - hparams.nextn_predict_layers);
 
                 switch (hparams.n_layer) {
                     case 20: model.type = MODEL_16B_A1B; break;
@@ -1411,6 +1411,71 @@ void llm_load_hparams(
                     default: model.type = e_model::MODEL_UNKNOWN;
                 }
             } break;
+        case LLM_ARCH_BAILINGMOE3:
+            {
+                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,       hparams.f_norm_rms_eps);
+                ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,         hparams.n_layer_dense_lead);
+                ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp);
+                ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
+                ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,               hparams.n_expert_shared);
+                ml.get_key(LLM_KV_EXPERT_GROUP_COUNT,                hparams.n_expert_groups);
+                ml.get_key(LLM_KV_EXPERT_GROUP_USED_COUNT,           hparams.n_group_used);
+                ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,              hparams.expert_weights_scale);
+                ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,               hparams.expert_weights_norm);
+                ml.get_key(LLM_KV_EXPERT_GATING_FUNC,                hparams.expert_gating_func);
+                // Ling-3.0-tiny ships no NextN block, and its converters omit the key
+                ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS,              hparams.nextn_predict_layers, false);
+                ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK,            hparams.n_lora_kv);
+                // Ling-3.0-flash sets q_lora_rank null and projects Q directly; Ling-3.0-tiny factorizes it
+                ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,             hparams.n_lora_q, false);
+                ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA,           hparams.n_embd_head_k_full);
+                ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA,         hparams.n_embd_head_v_full);
+                ml.get_key(LLM_KV_SSM_CONV_KERNEL,                    hparams.ssm_d_conv);
+                ml.get_key(LLM_KV_KDA_HEAD_DIM,                       hparams.ssm_d_state);
+                // absent key means true: only the safe-gate formula is implemented
+                hparams.kda_safe_gate = true;
+                ml.get_key(LLM_KV_KDA_SAFE_GATE,                      hparams.kda_safe_gate, false);
+                if (!hparams.kda_safe_gate) {
+                    throw std::runtime_error("bailingmoe3: kda.safe_gate = false is not supported");
+                }
+                ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND,               hparams.kda_gate_lower_bound);
+                // Ling-3.0-tiny sets both limit lists null. 0 is the unclamped value where it is read.
+                hparams.swiglu_limits.fill(0.0f);
+                hparams.swiglu_limits_shared.fill(0.0f);
+                ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,            hparams.swiglu_limits,        hparams.n_layer, false);
+                ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP,          hparams.swiglu_limits_shared, hparams.n_layer, false);
+
+                // one converter writes the tensors but not the key, so the tensors decide.
+                // the norm is 1-D and so cannot be transposed, unlike attn_q_a itself
+                if (hparams.n_lora_q == 0) {
+                    for (uint32_t il = 0; il < hparams.n_layer; ++il) {
+                        const std::string probe = LLM_TN(LLM_ARCH_BAILINGMOE3)(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il);
+                        if (const auto * meta = ml.get_tensor_meta(probe.c_str())) {
+                            hparams.n_lora_q = meta->ne[0];
+                            break;
+                        }
+                    }
+                }
+
+                hparams.ssm_n_group = hparams.n_head();
+                hparams.ssm_dt_rank = hparams.n_head();
+                hparams.ssm_d_inner = hparams.ssm_d_state * hparams.ssm_dt_rank;
+                // believe the tensors: a stale count hides the last real layer
+                if (hparams.nextn_predict_layers > 0) {
+                    const std::string probe = LLM_TN(LLM_ARCH_BAILINGMOE3)(LLM_TENSOR_NEXTN_EH_PROJ, "weight",
+                            hparams.n_layer - hparams.nextn_predict_layers);
+                    if (ml.get_tensor_meta(probe.c_str()) == nullptr) {
+                        hparams.nextn_predict_layers = 0;
+                    }
+                }
+                hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
+
+                for (uint32_t il = 0; il < hparams.n_layer; ++il) {
+                    hparams.recurrent_layer_arr[il] = hparams.n_head_kv_arr[il] == 0;
+                }
+
+                model.type = e_model::MODEL_UNKNOWN;
+            } break;
 	    case LLM_ARCH_DOTS1:
             {
                 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

src/llama-hparams.h

@@ -94,6 +94,8 @@ struct llama_hparams {
     uint32_t ssm_d_state = 0;
     uint32_t ssm_dt_rank = 0;
     uint32_t ssm_n_group = 0;
+    bool     kda_safe_gate = false;
+    float    kda_gate_lower_bound = 0.0f;
 
     // for hybrid state-space models (e.g. qwen3next)
     std::array<bool, LLAMA_MAX_LAYERS> recurrent_layer_arr;
@@ -226,6 +228,7 @@ struct llama_hparams {
         if (this->ssm_d_state != other.ssm_d_state) return true;
         if (this->ssm_dt_rank != other.ssm_dt_rank) return true;
         if (this->ssm_n_group != other.ssm_n_group) return true;
+        if (this->kda_safe_gate != other.kda_safe_gate) return true;
         if (this->recurrent_layer_arr != other.recurrent_layer_arr) return true;
         for (int i = 0; i < 8; ++i) {
             if (this->dflash_target_layer_ids[i] != other.dflash_target_layer_ids[i]) return true;
@@ -237,6 +240,7 @@ struct llama_hparams {
 
         if (!is_float_close(this->f_norm_eps,            other.f_norm_eps,            EPSILON)) return true;
         if (!is_float_close(this->f_norm_rms_eps,        other.f_norm_rms_eps,        EPSILON)) return true;
+        if (!is_float_close(this->kda_gate_lower_bound,  other.kda_gate_lower_bound,  EPSILON)) return true;
         if (!is_float_close(this->rope_attn_factor,      other.rope_attn_factor,      EPSILON)) return true;
         if (!is_float_close(this->rope_freq_base_train,  other.rope_freq_base_train,  EPSILON)) return true;
         if (!is_float_close(this->rope_freq_scale_train, other.rope_freq_scale_train, EPSILON)) return true;

src/llama-kda.cpp

@@ -0,0 +1,250 @@
+#include "llama-kda.h"
+#include "llama-hparams.h"
+#include "llama-cparams.h"
+#include "llama-model.h"
+#include "llama-context.h"
+
+#include "ggml.h"
+
+static std::pair<ggml_tensor *, ggml_tensor *> build_kda_qkvz(llama_context & lctx, ggml_context * ctx0,
+        ggml_tensor * wq, ggml_tensor * wk, ggml_tensor * wv, ggml_tensor * ssm_g_a,
+        ggml_tensor * input, int il, const llm_build_cb & cb, ggml_cgraph * gf) {
+    auto q = llm_build_context::llm_build_lora_mm(lctx, ctx0, wq, input);
+    auto k = llm_build_context::llm_build_lora_mm(lctx, ctx0, wk, input);
+    auto v = llm_build_context::llm_build_lora_mm(lctx, ctx0, wv, input);
+    auto z = llm_build_context::llm_build_lora_mm(lctx, ctx0, ssm_g_a, input);
+    cb(q, "q", il);
+    cb(k, "k", il);
+    cb(v, "v", il);
+    cb(z, "z", il);
+
+    auto qkv = ggml_concat(ctx0, q, k, 0);
+    qkv = ggml_concat(ctx0, qkv, v, 0);
+    cb(qkv, "qkv_mixed", il);
+    ggml_build_forward_expand(gf, qkv);
+    ggml_build_forward_expand(gf, z);
+    return { qkv, z };
+}
+
+static std::pair<ggml_tensor *, ggml_tensor *> build_kda_beta_gate(llama_context & lctx, ggml_context * ctx0,
+        ggml_tensor * ssm_beta, ggml_tensor * ssm_f_a, ggml_tensor * ssm_dt_b, ggml_tensor * ssm_a,
+        ggml_tensor * input, int64_t head_dim, int64_t n_head, float lower_bound,
+        int il, const llm_build_cb & cb, ggml_cgraph * gf) {
+    const int64_t n_tok = input->ne[1];
+
+    auto beta = llm_build_context::llm_build_lora_mm(lctx, ctx0, ssm_beta, input);
+    beta = ggml_reshape_4d(ctx0, beta, n_head, 1, n_tok, 1);
+    cb(beta, "beta", il);
+
+    auto raw = llm_build_context::llm_build_lora_mm(lctx, ctx0, ssm_f_a, input);
+    raw = ggml_reshape_4d(ctx0, raw, head_dim, n_head, n_tok, 1);
+    cb(raw, "decay_raw", il);
+
+    auto dt = ggml_reshape_4d(ctx0, ssm_dt_b, head_dim, n_head, 1, 1);
+    auto a  = ggml_reshape_4d(ctx0, ssm_a, 1, n_head, 1, 1);
+    auto log_decay = ggml_add(ctx0, raw, dt);
+    log_decay = ggml_mul(ctx0, log_decay, a);
+    log_decay = ggml_sigmoid(ctx0, log_decay);
+    log_decay = ggml_scale(ctx0, log_decay, lower_bound);
+    cb(log_decay, "log_decay", il);
+
+    ggml_build_forward_expand(gf, beta);
+    ggml_build_forward_expand(gf, log_decay);
+    return { beta, log_decay };
+}
+
+static ggml_tensor * build_kda_conv(ggml_context * ctx0,
+        ggml_tensor * ssm_conv1d_q, ggml_tensor * ssm_conv1d_k, ggml_tensor * ssm_conv1d_v) {
+    const int64_t d_conv = ssm_conv1d_q->ne[0];
+    auto q = ggml_reshape_2d(ctx0, ssm_conv1d_q, d_conv, ggml_nrows(ssm_conv1d_q));
+    auto k = ggml_reshape_2d(ctx0, ssm_conv1d_k, d_conv, ggml_nrows(ssm_conv1d_k));
+    auto v = ggml_reshape_2d(ctx0, ssm_conv1d_v, d_conv, ggml_nrows(ssm_conv1d_v));
+    auto qkv = ggml_concat(ctx0, q, k, 1);
+    return ggml_concat(ctx0, qkv, v, 1);
+}
+
+static ggml_tensor * build_kda_gated_output(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_norm, ggml_tensor * ssm_out, ggml_tensor * output, ggml_tensor * z,
+        int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb) {
+
+    ggml_tensor * attn_out_2d = ggml_reshape_2d(ctx0, output, head_v_dim, num_v_heads * n_tok);
+    ggml_tensor * z_2d        = ggml_reshape_2d(ctx0, z,      head_v_dim, num_v_heads * n_tok);
+
+    ggml_tensor * attn_out_norm = llm_build_context::llm_build_norm(ctx0, attn_out_2d, lctx.model.hparams, ssm_norm, nullptr, LLM_NORM_RMS, cb, il);
+    cb(attn_out_norm, "attn_rms_norm", il);
+    attn_out_norm = ggml_mul(ctx0, ggml_sigmoid(ctx0, z_2d), attn_out_norm);
+    cb(attn_out_norm, "attn_out_norm", il);
+
+    ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim*num_v_heads, n_tok);
+    cb(final_output, "final_output", il);
+
+    ggml_tensor * out = llm_build_context::llm_build_lora_mm(lctx, ctx0, ssm_out, final_output);
+    cb(out, "linear_attn_out", il);
+
+    return ggml_reshape_2d(ctx0, out, lctx.model.hparams.n_embd, n_tok);
+}
+
+ggml_tensor * delta_net::build_layer_attn_kda_core(ggml_context * ctx0, ggml_cgraph * gf,
+        ggml_tensor * delta_input, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids,
+        uint32_t state_seq_id_local, bool reset_state_local, int il, const llm_build_cb & cb) const {
+    const int64_t n_tok = delta_input->ne[1];
+    const int64_t head_dim = lctx.model.hparams.ssm_d_state;
+
+    auto & model   = lctx.model;
+    auto & hparams = model.hparams;
+    auto & kv_self = lctx.kv_self;
+    auto & layer   = model.layers[il];
+
+    if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH && kv_self.s_l[il]->extra) {
+        auto split_s_l = (ggml_split_tensor_t *) kv_self.s_l[il]->extra;
+        GGML_ASSERT(split_s_l && split_s_l->n_device > 1);
+
+        std::vector<ggml_tensor *> results(split_s_l->n_device, nullptr);
+        bool input_added = false;
+        for (int id = 0; id < split_s_l->n_device; ++id) {
+            if (!split_s_l->splits[id]) {
+                continue;
+            }
+
+            auto split = [id](ggml_tensor * tensor) {
+                auto data = (ggml_split_tensor_t *) tensor->extra;
+                GGML_ASSERT(data && data->splits[id]);
+                return data->splits[id];
+            };
+
+            const int il_cb = 1000 * il + id;
+            auto input = llm_build_context::get_input_tensor_sm_graph(ctx0, delta_input, id);
+            auto cur = llm_build_context::llm_build_norm(ctx0, input, hparams,
+                    split(layer.attn_norm), nullptr, LLM_NORM_RMS, cb, il_cb);
+
+            auto ssm_out = split(layer.ssm_out);
+            const int64_t n_head = ssm_out->ne[0] / head_dim;
+            auto [qkv_mixed, z] = build_kda_qkvz(lctx, ctx0,
+                    split(layer.wq), split(layer.wk), split(layer.wv), split(layer.ssm_g_a),
+                    cur, il_cb, cb, gf);
+            auto [beta, log_decay] = build_kda_beta_gate(lctx, ctx0,
+                    split(layer.ssm_beta), split(layer.ssm_f_a), split(layer.ssm_dt_b), split(layer.ssm_a),
+                    cur, head_dim, n_head, hparams.kda_gate_lower_bound, il_cb, cb, gf);
+            auto conv = build_kda_conv(ctx0,
+                    split(layer.ssm_conv1d_q), split(layer.ssm_conv1d_k), split(layer.ssm_conv1d_v));
+
+            ggml_tensor * per_step_ckpt = nullptr;
+            if (save_per_step_states && il < (int) kv_self.ckpt.per_step_ssm.size()) {
+                per_step_ckpt = kv_self.ckpt.per_step_ssm[il][id];
+            }
+            auto per_step_conv = save_per_step_states && il < (int) kv_self.ckpt.per_step_conv.size() &&
+                                 id < (int) kv_self.ckpt.per_step_conv[il].size()
+                               ? kv_self.ckpt.per_step_conv[il][id] : nullptr;
+
+            const uint32_t qnext_state_slots = split_s_l->splits[id]->ne[1];
+            auto output = build_qkv(ctx0, split_s_l->splits[id], conv, qkv_mixed,
+                    inp_s_seq_qnext, beta, log_decay,
+                    head_dim, n_head, head_dim, n_head, hparams.ssm_d_conv,
+                    state_seq_id_local, qnext_state_slots, reset_state_local,
+                    hparams.f_norm_rms_eps, 1, il_cb, cb, gf, per_step_ckpt, per_step_conv);
+
+            auto gated_output = build_kda_gated_output(lctx, ctx0, split(layer.ssm_norm), ssm_out, output, z,
+                    head_dim, n_head, n_tok, il_cb, cb);
+            if (inp_out_ids) {
+                gated_output = ggml_get_rows(ctx0, gated_output, inp_out_ids);
+            }
+            if (!input_added) {
+                if (inp_out_ids) {
+                    input = ggml_get_rows(ctx0, input, inp_out_ids);
+                }
+                gated_output = ggml_add(ctx0, gated_output, input);
+                input_added = true;
+            }
+            if (gated_output->ne[1] > 32 && lctx.cparams.reduce_type != GGML_TYPE_F32) {
+                gated_output = ggml_cast(ctx0, gated_output, lctx.cparams.reduce_type);
+            }
+            ggml_build_forward_expand(gf, gated_output);
+            results[id] = gated_output;
+        }
+
+        auto output = ggml_reduce(ctx0, results.data(), split_s_l->n_device, GGML_OP_ADD);
+        ggml_build_forward_expand(gf, output);
+        return output;
+    }
+
+    const uint32_t qnext_state_slots = llm_build_context::llama_kv_qnext_state_slots(kv_self);
+    int idx = model.default_layer_device[il];
+    auto input = delta_input;
+    if (input->op == GGML_OP_REDUCE) {
+        const int idx_s_l = ggml_backend_sched_get_backend_idx(lctx.sched, kv_self.s_l[il]->buffer);
+        if (idx_s_l >= 0) {
+            idx = idx_s_l;
+        }
+        if (input->src[idx]) {
+            input->view_src = input->src[idx];
+        }
+    }
+
+    auto norm = layer.attn_norm->extra
+              ? ((ggml_split_tensor_t *) layer.attn_norm->extra)->splits[idx]
+              : layer.attn_norm;
+    auto cur = llm_build_context::llm_build_norm(ctx0, input, hparams, norm, nullptr, LLM_NORM_RMS, cb, il);
+
+    const int64_t n_head = hparams.ssm_dt_rank;
+    auto [qkv_mixed, z] = build_kda_qkvz(lctx, ctx0,
+            layer.wq, layer.wk, layer.wv, layer.ssm_g_a, cur, il, cb, gf);
+    auto [beta, log_decay] = build_kda_beta_gate(lctx, ctx0,
+            layer.ssm_beta, layer.ssm_f_a, layer.ssm_dt_b, layer.ssm_a,
+            cur, head_dim, n_head, hparams.kda_gate_lower_bound, il, cb, gf);
+    auto conv = build_kda_conv(ctx0, layer.ssm_conv1d_q, layer.ssm_conv1d_k, layer.ssm_conv1d_v);
+
+    ggml_tensor * per_step_ckpt = nullptr;
+    if (save_per_step_states && il < (int) kv_self.ckpt.per_step_ssm.size()) {
+        per_step_ckpt = kv_self.ckpt.per_step_ssm[il].front();
+    }
+    auto per_step_conv = save_per_step_states && il < (int) kv_self.ckpt.per_step_conv.size() &&
+                         !kv_self.ckpt.per_step_conv[il].empty()
+                       ? kv_self.ckpt.per_step_conv[il].front() : nullptr;
+
+    auto output = build_qkv(ctx0, kv_self.s_l[il], conv, qkv_mixed,
+            inp_s_seq_qnext, beta, log_decay,
+            head_dim, n_head, head_dim, n_head, hparams.ssm_d_conv,
+            state_seq_id_local, qnext_state_slots, reset_state_local,
+            hparams.f_norm_rms_eps, 1, il, cb, gf, per_step_ckpt, per_step_conv);
+    auto gated_output = build_kda_gated_output(lctx, ctx0, layer.ssm_norm, layer.ssm_out, output, z,
+            head_dim, n_head, n_tok, il, cb);
+
+    if (inp_out_ids) {
+        gated_output = ggml_get_rows(ctx0, gated_output, inp_out_ids);
+        input = ggml_get_rows(ctx0, input, inp_out_ids);
+    }
+    output = ggml_add(ctx0, gated_output, input);
+    cb(output, "ssm_output", il);
+    return output;
+}
+
+ggml_tensor * delta_net::build_layer_attn_kda(ggml_context * ctx0, ggml_cgraph * gf,
+        ggml_tensor * cur, ggml_tensor * inp_out_ids, int il, const llm_build_cb & cb) const {
+    GGML_ASSERT(lctx.inp_s_seq_qnext != nullptr);
+
+    auto & layer = lctx.model.layers[il];
+    GGML_ASSERT(lctx.model.hparams.is_recurrent(il));
+    GGML_ASSERT(layer.wq && layer.wk && layer.wv);
+    GGML_ASSERT(layer.ssm_conv1d_q && layer.ssm_conv1d_k && layer.ssm_conv1d_v);
+    GGML_ASSERT(layer.ssm_f_a && layer.ssm_g_a && layer.ssm_beta);
+    GGML_ASSERT(layer.ssm_a && layer.ssm_dt_b && layer.ssm_norm && layer.ssm_out);
+
+    if (all_same_seq) {
+        const bool reset_state = batch.pos != nullptr && batch.pos[0] == 0;
+        return build_layer_attn_kda_core(ctx0, gf, cur, lctx.inp_s_seq_qnext, inp_out_ids,
+                token_seq_ids.front(), reset_state, il, cb);
+    }
+
+    GGML_ASSERT(has_unique_seq_ids && "bailingmoe3 mixed-sequence batches require unique sequence IDs per token");
+
+    ggml_tensor * out = nullptr;
+    for (int64_t i = 0; i < batch.n_tokens; ++i) {
+        auto cur_i = ggml_view_2d(ctx0, cur, cur->ne[0], 1, cur->nb[1], (size_t) i * cur->nb[1]);
+        auto inp_s_seq_qnext_i = ggml_view_2d(ctx0, lctx.inp_s_seq_qnext, 1, 1,
+                lctx.inp_s_seq_qnext->nb[1], (size_t) i * lctx.inp_s_seq_qnext->nb[1]);
+        const bool reset_state = batch.pos != nullptr && batch.pos[i] == 0;
+        auto out_i = build_layer_attn_kda_core(ctx0, gf, cur_i, inp_s_seq_qnext_i, inp_out_ids,
+                (uint32_t) token_seq_ids[i], reset_state, il, cb);
+        out = out == nullptr ? out_i : ggml_concat(ctx0, out, out_i, 1);
+    }
+    return out;
+}

src/llama-kda.h

@@ -0,0 +1,3 @@
+#pragma once
+
+#include "llama-delta-net.h"

src/llama-load-tensors.cpp

@@ -157,6 +157,7 @@ struct create_tensors_helper : public create_tensors_helper_interface {
     bool create_openai_moe_tensors(const LLM_TN & tn);
 
     bool create_bailingmoe2_tensors(const LLM_TN & tn);
+    bool create_bailingmoe3_tensors(const LLM_TN & tn);
 
     bool create_minimaxm2_tensors(const LLM_TN & tn);
     bool create_minimaxm3_tensors(const LLM_TN & tn);
@@ -4049,6 +4050,107 @@ bool create_tensors_helper::create_bailingmoe2_tensors(const LLM_TN & tn) {
     return use_mmap_buffer;
 }
 
+bool create_tensors_helper::create_bailingmoe3_tensors(const LLM_TN & tn) {
+    LOADING_PRELUDE
+
+    const uint32_t n_head_kda       = hparams.ssm_dt_rank;
+    const uint32_t n_embd_kda       = hparams.ssm_d_inner;
+    const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot;
+    const uint32_t n_ff_exp         = hparams.n_ff_exp;
+    const uint32_t n_ff_shexp       = hparams.n_ff_shexp * hparams.n_expert_shared;
+
+    model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab});
+    model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd});
+    model.output      = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab});
+
+    for (uint32_t il = 0; il < hparams.n_layer; ++il) {
+        auto & layer = model.layers[il];
+        ggml_context * ctx_layer = ctx_for_layer(il);
+        ggml_context * ctx_split = ctx_for_layer_split(il);
+        const bool graph_or_attn = model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN;
+        ggml_context * norm_ctx = graph_or_attn ? ctx_split : ctx_layer;
+        ggml_context * moe_ctx  = graph_or_attn ? ctx_split : ctx_layer;
+        const bool is_mtp_layer = !hparams.has_kv(il);
+        const int flags = is_mtp_layer ? llama_model_loader::TENSOR_SKIP : 0;
+        const bool is_kda = hparams.is_recurrent(il);
+
+        layer.attn_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_NORM, "weight", il), {n_embd}, flags);
+        if (is_kda) {
+            layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", il), {n_embd, n_embd_kda}, flags);
+            layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", il), {n_embd, n_embd_kda}, flags);
+            layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", il), {n_embd, n_embd_kda}, flags);
+            layer.ssm_conv1d_q = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il),
+                    {hparams.ssm_d_conv, 1, n_embd_kda}, flags);
+            layer.ssm_conv1d_k = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il),
+                    {hparams.ssm_d_conv, 1, n_embd_kda}, flags);
+            layer.ssm_conv1d_v = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il),
+                    {hparams.ssm_d_conv, 1, n_embd_kda}, flags);
+            // no LoRA factorization here, so converters disagree on the _a suffix
+            layer.ssm_f_a  = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_F_A,  "weight", il), {n_embd, n_embd_kda},
+                    flags | llama_model_loader::TENSOR_NOT_REQUIRED);
+            if (!layer.ssm_f_a) {
+                layer.ssm_f_a = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_F, "weight", il), {n_embd, n_embd_kda}, flags);
+            }
+            layer.ssm_g_a  = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_G_A,  "weight", il), {n_embd, n_embd_kda},
+                    flags | llama_model_loader::TENSOR_NOT_REQUIRED);
+            if (!layer.ssm_g_a) {
+                layer.ssm_g_a = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_G, "weight", il), {n_embd, n_embd_kda}, flags);
+            }
+            layer.ssm_beta = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_BETA, "weight", il), {n_embd, n_head_kda}, flags);
+            layer.ssm_a    = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_A,              il), {1, n_head_kda}, flags);
+            layer.ssm_dt_b = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_DT,   "bias",   il), {n_embd_kda}, flags);
+            layer.ssm_norm = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_NORM, "weight", il), {hparams.ssm_d_state}, flags);
+            layer.ssm_out  = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", il), {n_embd_kda, n_embd}, flags);
+        } else {
+            if (hparams.n_lora_q > 0) {
+                layer.wq_a = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_A, "weight", il),
+                        {n_embd, hparams.n_lora_q}, flags);
+                layer.attn_q_a_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il),
+                        {hparams.n_lora_q}, flags);
+                layer.wq_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_B, "weight", il),
+                        {hparams.n_lora_q, n_head_kda * hparams.n_embd_head_k_full}, flags);
+            } else {
+                layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", il),
+                        {n_embd, n_head_kda * hparams.n_embd_head_k_full}, flags);
+            }
+            layer.wkv_a_mqa = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il),
+                    {n_embd, hparams.n_lora_kv + hparams.n_rot}, flags);
+            layer.attn_kv_a_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il),
+                    {hparams.n_lora_kv}, flags);
+            layer.wk_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_B, "weight", il),
+                    {n_embd_head_qk_nope, hparams.n_lora_kv, n_head_kda}, flags);
+            layer.wv_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V_B, "weight", il),
+                    {hparams.n_lora_kv, hparams.n_embd_head_v(0), n_head_kda}, flags);
+            layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", il),
+                    {n_embd, n_head_kda}, flags);
+            layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", il), {n_embd_kda, n_embd}, flags);
+        }
+
+        layer.ffn_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_FFN_NORM, "weight", il), {n_embd}, flags);
+        if (il < hparams.n_layer_dense_lead) {
+            layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, flags);
+            layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", il), {n_ff, n_embd}, flags);
+            layer.ffn_up   = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP,   "weight", il), {n_embd, n_ff}, flags);
+        } else {
+            layer.ffn_gate_inp = create_tensor(moe_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), {n_embd, n_expert}, flags);
+            layer.ffn_exp_probs_b = create_tensor(moe_ctx, tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), {n_expert}, flags);
+            use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, il, flags, n_ff_exp, ctx_split);
+            layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), {n_embd, n_ff_shexp}, flags);
+            layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), {n_ff_shexp, n_embd}, flags);
+            layer.ffn_up_shexp   = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", il), {n_embd, n_ff_shexp}, flags);
+        }
+
+        if (is_mtp_layer) {
+            layer.nextn.eh_proj = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), {2 * n_embd, n_embd}, flags);
+            layer.nextn.enorm   = create_tensor(ctx_layer, tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), {n_embd}, flags);
+            layer.nextn.hnorm   = create_tensor(ctx_layer, tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), {n_embd}, flags);
+            layer.layer_out_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", il), {n_embd}, flags);
+        }
+    }
+
+    return use_mmap_buffer;
+}
+
 bool create_tensors_helper::create_ernie45_tensors(const LLM_TN & tn) {
     LOADING_PRELUDE
 
@@ -4669,12 +4771,17 @@ static void distribute_mla_tensors_for_split_mode_graph(
     if (layer.wq_b) {
         prepare_split_tensors(1, ctx_split, layer.wq_b, layer.split_wq_b, split_wq_b_cols, mem_used);
     } else if (layer.wq) {
-        // DSV2-Lite / no-Q-LoRA path: column-split wq directly along the head dim.
+        // serialized head width can differ from the MLA cache head width
         auto split_wq_cols = split_heads;
-        for (auto & s : split_wq_cols) s *= n_embd_head_k;
+        const int n_embd_head_q = layer.wq->ne[1] / n_head;
+        for (auto & s : split_wq_cols) s *= n_embd_head_q;
         prepare_split_tensors(1, ctx_split, layer.wq, layer.split_wq, split_wq_cols, mem_used);
     }
 
+    if (layer.wqkv_gate) {
+        prepare_split_tensors(1, ctx_split, layer.wqkv_gate, layer.split_wqkv_gate, split_heads, mem_used);
+    }
+
     // wkv_a_mqa is replicated (mirror): its per-head 3D batched mul_mat can't read a split src0.
     // wk_b/wv_b are split per head (split_dim=2); reshape a 2D folded-head GGUF layout to 3D first.
     if (layer.wkv_a_mqa) {
@@ -4995,6 +5102,38 @@ static void split_recurrent_tensors(const llama_hparams & hparams, llama_layer &
     LLAMA_LOG_DEBUG("    original size: %g MiB, split size: %g MiB\n", orig_size/1024./1024., split_size/1024./1024.);
 }
 
+static void split_bailingmoe3_kda_tensors(const llama_hparams & hparams, llama_layer & layer,
+        const std::vector<float> & cur_splits, std::vector<size_t> & mem_used, ggml_context * ctx_split) {
+    const int n_head = hparams.ssm_dt_rank;
+    const int n_embd_head = hparams.ssm_d_state;
+
+    int head_granularity = 1;
+    const auto type_traits = ggml_internal_get_type_traits(layer.ssm_out->type);
+    if (type_traits.blck_size > n_embd_head) {
+        head_granularity = type_traits.blck_size / n_embd_head;
+    }
+
+    auto split_heads = create_split(n_head, head_granularity, cur_splits, mem_used);
+    auto split_channels = split_heads;
+    for (auto & channels : split_channels) {
+        channels *= n_embd_head;
+    }
+
+    prepare_split_tensors( 1, ctx_split, layer.wq,           layer.split_wq,           split_channels, mem_used);
+    prepare_split_tensors( 1, ctx_split, layer.wk,           layer.split_wk,           split_channels, mem_used);
+    prepare_split_tensors( 1, ctx_split, layer.wv,           layer.split_wv,           split_channels, mem_used);
+    prepare_split_tensors( 2, ctx_split, layer.ssm_conv1d_q, layer.split_ssm_conv1d_q, split_channels, mem_used);
+    prepare_split_tensors( 2, ctx_split, layer.ssm_conv1d_k, layer.split_ssm_conv1d_k, split_channels, mem_used);
+    prepare_split_tensors( 2, ctx_split, layer.ssm_conv1d_v, layer.split_ssm_conv1d_v, split_channels, mem_used);
+    prepare_split_tensors( 1, ctx_split, layer.ssm_f_a,      layer.split_ssm_f_a,      split_channels, mem_used);
+    prepare_split_tensors( 1, ctx_split, layer.ssm_g_a,      layer.split_ssm_g_a,      split_channels, mem_used);
+    prepare_split_tensors( 1, ctx_split, layer.ssm_beta,     layer.split_ssm_beta,     split_heads,    mem_used);
+    prepare_split_tensors( 1, ctx_split, layer.ssm_a,        layer.split_ssm_a,        split_heads,    mem_used);
+    prepare_split_tensors( 0, ctx_split, layer.ssm_dt_b,     layer.split_ssm_dt,       split_channels, mem_used);
+    prepare_split_tensors(-1, ctx_split, layer.ssm_norm,     layer.split_ssm_norm,     split_heads,    mem_used);
+    prepare_split_tensors( 0, ctx_split, layer.ssm_out,      layer.split_ssm_out,      split_channels, mem_used);
+}
+
 bool create_tensors_helper::create_tensors() {
     const auto tn = LLM_TN(model.arch);
     bool use_mmap_buffer = true;
@@ -5147,6 +5286,8 @@ bool create_tensors_helper::create_tensors() {
             use_mmap_buffer = create_openai_moe_tensors(tn); break;
         case LLM_ARCH_BAILINGMOE2:
             use_mmap_buffer = create_bailingmoe2_tensors(tn); break;
+        case LLM_ARCH_BAILINGMOE3:
+            use_mmap_buffer = create_bailingmoe3_tensors(tn); break;
         case LLM_ARCH_MINIMAX_M2:
             use_mmap_buffer = create_minimaxm2_tensors(tn); break;
         case LLM_ARCH_MINIMAX_M3:
@@ -5229,7 +5370,6 @@ bool create_tensors_helper::create_tensors() {
                 LLAMA_LOG_DEBUG("%s: not splitting MTP tail layer %d (forced non-split)\n", __func__, il);
                 continue;
             }
-            int gqa_ratio = hparams.n_head(il) / hparams.n_head_kv(il);
             if (ggml_backend_buft_is_host(model.buft_layer[il].buft_matrix)) {
                 LLAMA_LOG_INFO("%s: not splitting layer %d because buffer type is host\n", __func__, il);
                 continue;
@@ -5267,7 +5407,11 @@ bool create_tensors_helper::create_tensors() {
                 prepare_split_tensors(-1, ctx_split, layer.rope_freqs, layer.split_rope_freqs, split, mem_used);
             }
             if (hparams.is_recurrent(il)) {
-                split_recurrent_tensors(hparams, layer, cur_splits, mem_used, ctx_split, il); //, model.arch == LLM_ARCH_QWEN3NEXT ? 0 : 1);
+                if (model.arch == LLM_ARCH_BAILINGMOE3) {
+                    split_bailingmoe3_kda_tensors(hparams, layer, cur_splits, mem_used, ctx_split);
+                } else {
+                    split_recurrent_tensors(hparams, layer, cur_splits, mem_used, ctx_split, il);
+                }
             }
             else if (is_gemma4_assistant()) {
                 GGML_ASSERT(layer.wo && layer.wq);
@@ -5312,6 +5456,7 @@ bool create_tensors_helper::create_tensors() {
                 }
             }
             else if (layer.wo && layer.wq && layer.wk && (layer.wv || model.arch == LLM_ARCH_GEMMA4)) {
+                const int gqa_ratio = hparams.n_head(il) / hparams.n_head_kv(il);
                 auto granularity_kq = hparams.n_embd_head_k(il) * gqa_ratio;
                 int wq_ne1 = layer.wq->ne[1];
                 if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN35) {
@@ -5464,12 +5609,8 @@ bool create_tensors_helper::create_tensors() {
                 }
             }
 
-            // MLA tensor distribution (DEEPSEEK2/GLM_DSA/MISTRAL4). Detect by arch + absence of wk
-            // since wkv_b can be null when the model was quantized by mainline llama.cpp.
-            if (layer.wo && !layer.wk &&
-                (model.arch == LLM_ARCH_DEEPSEEK2 ||
-                 model.arch == LLM_ARCH_GLM_DSA ||
-                 model.arch == LLM_ARCH_MISTRAL4)) {
+            // wkv_b can be null when the model was quantized by mainline llama.cpp.
+            if (layer.wo && !layer.wk && model.is_mla_model()) {
                 distribute_mla_tensors_for_split_mode_graph(
                     layer, hparams, cur_splits, mem_used, ctx_split, il);
             }

src/llama-model.cpp

@@ -1663,6 +1663,55 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
             { LLM_TENSOR_LAYER_OUT_NORM,     "blk.%d.layer_output_norm" },
         },
     },
+    {
+        LLM_ARCH_BAILINGMOE3,
+        {
+            { LLM_TENSOR_TOKEN_EMBD,         "token_embd" },
+            { LLM_TENSOR_OUTPUT_NORM,        "output_norm" },
+            { LLM_TENSOR_OUTPUT,             "output" },
+            { LLM_TENSOR_ATTN_NORM,          "blk.%d.attn_norm" },
+            { LLM_TENSOR_ATTN_Q,             "blk.%d.attn_q" },
+            { LLM_TENSOR_ATTN_K,             "blk.%d.attn_k" },
+            { LLM_TENSOR_ATTN_V,             "blk.%d.attn_v" },
+            { LLM_TENSOR_ATTN_OUT,           "blk.%d.attn_output" },
+            { LLM_TENSOR_ATTN_GATE,          "blk.%d.attn_gate" },
+            { LLM_TENSOR_ATTN_Q_A,           "blk.%d.attn_q_a" },
+            { LLM_TENSOR_ATTN_Q_A_NORM,      "blk.%d.attn_q_a_norm" },
+            { LLM_TENSOR_ATTN_Q_B,           "blk.%d.attn_q_b" },
+            { LLM_TENSOR_ATTN_KV_A_MQA,      "blk.%d.attn_kv_a_mqa" },
+            { LLM_TENSOR_ATTN_KV_A_NORM,     "blk.%d.attn_kv_a_norm" },
+            { LLM_TENSOR_ATTN_K_B,           "blk.%d.attn_k_b" },
+            { LLM_TENSOR_ATTN_V_B,           "blk.%d.attn_v_b" },
+            { LLM_TENSOR_SSM_CONV1D_Q,       "blk.%d.ssm_conv1d_q" },
+            { LLM_TENSOR_SSM_CONV1D_K,       "blk.%d.ssm_conv1d_k" },
+            { LLM_TENSOR_SSM_CONV1D_V,       "blk.%d.ssm_conv1d_v" },
+            { LLM_TENSOR_SSM_F_A,            "blk.%d.ssm_f_a" },
+            { LLM_TENSOR_SSM_G_A,            "blk.%d.ssm_g_a" },
+            { LLM_TENSOR_SSM_F,              "blk.%d.ssm_f" },
+            { LLM_TENSOR_SSM_G,              "blk.%d.ssm_g" },
+            { LLM_TENSOR_SSM_A,              "blk.%d.ssm_a" },
+            { LLM_TENSOR_SSM_DT,             "blk.%d.ssm_dt" },
+            { LLM_TENSOR_SSM_BETA,           "blk.%d.ssm_beta" },
+            { LLM_TENSOR_SSM_NORM,           "blk.%d.ssm_norm" },
+            { LLM_TENSOR_FFN_NORM,           "blk.%d.ffn_norm" },
+            { LLM_TENSOR_FFN_GATE,           "blk.%d.ffn_gate" },
+            { LLM_TENSOR_FFN_DOWN,           "blk.%d.ffn_down" },
+            { LLM_TENSOR_FFN_UP,             "blk.%d.ffn_up" },
+            { LLM_TENSOR_FFN_GATE_INP,       "blk.%d.ffn_gate_inp" },
+            { LLM_TENSOR_FFN_EXP_PROBS_B,    "blk.%d.exp_probs_b" },
+            { LLM_TENSOR_FFN_GATE_EXPS,      "blk.%d.ffn_gate_exps" },
+            { LLM_TENSOR_FFN_DOWN_EXPS,      "blk.%d.ffn_down_exps" },
+            { LLM_TENSOR_FFN_UP_EXPS,        "blk.%d.ffn_up_exps" },
+            { LLM_TENSOR_FFN_GATE_UP_EXPS,   "blk.%d.ffn_gate_up_exps" },
+            { LLM_TENSOR_FFN_GATE_SHEXP,     "blk.%d.ffn_gate_shexp" },
+            { LLM_TENSOR_FFN_DOWN_SHEXP,     "blk.%d.ffn_down_shexp" },
+            { LLM_TENSOR_FFN_UP_SHEXP,       "blk.%d.ffn_up_shexp" },
+            { LLM_TENSOR_NEXTN_EH_PROJ,      "blk.%d.nextn.eh_proj" },
+            { LLM_TENSOR_NEXTN_ENORM,        "blk.%d.nextn.enorm" },
+            { LLM_TENSOR_NEXTN_HNORM,        "blk.%d.nextn.hnorm" },
+            { LLM_TENSOR_LAYER_OUT_NORM,     "blk.%d.layer_output_norm" },
+        },
+    },
     {
         LLM_ARCH_MINIMAX_M2,
         {
@@ -2433,7 +2482,7 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_
         }
         return size;
     }
-    bool is_mla_attn = arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4;
+    bool is_mla_attn = is_mla_model();
     if (is_mla_attn && mla_attn) {
         auto n_embd_head_qk_rope = hparams.n_rot;
         auto kv_lora_rank = hparams.n_lora_kv;

src/llama-model.h

@@ -243,11 +243,16 @@ struct llama_layer {
     llama_split_tensor split_ssm_wqkv_gate;
     llama_split_tensor split_ssm_in;
     llama_split_tensor split_ssm_conv1d;
+    llama_split_tensor split_ssm_conv1d_q;
+    llama_split_tensor split_ssm_conv1d_k;
+    llama_split_tensor split_ssm_conv1d_v;
     llama_split_tensor split_ssm_dt;
     llama_split_tensor split_ssm_a;
     llama_split_tensor split_ssm_beta_alpha;
     llama_split_tensor split_ssm_beta;
     llama_split_tensor split_ssm_alpha;
+    llama_split_tensor split_ssm_f_a;
+    llama_split_tensor split_ssm_g_a;
     llama_split_tensor split_ssm_norm;
     llama_split_tensor split_ssm_out;
 
@@ -354,9 +359,14 @@ struct llama_layer {
     struct ggml_tensor * ssm_beta_alpha = nullptr;
     struct ggml_tensor * ssm_alpha = nullptr;
     struct ggml_tensor * ssm_beta = nullptr;
+    struct ggml_tensor * ssm_f_a = nullptr;
+    struct ggml_tensor * ssm_g_a = nullptr;
 
     // mamba
     struct ggml_tensor * ssm_conv1d = nullptr;
+    struct ggml_tensor * ssm_conv1d_q = nullptr;
+    struct ggml_tensor * ssm_conv1d_k = nullptr;
+    struct ggml_tensor * ssm_conv1d_v = nullptr;
     struct ggml_tensor * ssm_a = nullptr;
     struct ggml_tensor * ssm_d = nullptr;
 
@@ -581,7 +591,14 @@ struct llama_model {
     }
 
     bool is_mla_model() const {
-        return arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4;
+        return arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_BAILINGMOE3;
+    }
+
+    float swiglu_limit(uint32_t il, bool shared) const {
+        if (arch != LLM_ARCH_STEP35 && arch != LLM_ARCH_BAILINGMOE3 && arch != LLM_ARCH_DEEPSEEK4) {
+            return 0.0f;
+        }
+        return shared ? hparams.swiglu_limits_shared[il] : hparams.swiglu_limits[il];
     }
 
     // a compacted sliding-window cache needs the graph to build its KQ mask over the compacted

src/llama-quantize.cpp

@@ -1484,7 +1484,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
 
         // do not quantize Mamba's small yet 2D weights
         // NOTE: can't use LLM_TN here because the layer number is not known
-        quantize &= name.find("ssm_conv1d.weight") == std::string::npos;
+        quantize &= name.find("ssm_conv1d")        == std::string::npos;
         quantize &= name.find("ssm_x.weight")      == std::string::npos;
         quantize &= name.find("ssm_dt.weight")     == std::string::npos;
 

src/llama.cpp

@@ -1254,6 +1254,9 @@ static bool llama_kv_cache_init(
     }
     if ((model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN) && is_mla_attn && offload) {
         cache.replicated_k_l.reserve(n_layer);
+        if (llama_model_has_recurrent(&model)) {
+            cache.split_s_l.reserve(n_layer);
+        }
         replicate_mla = true;
     }
 
@@ -1377,6 +1380,42 @@ static bool llama_kv_cache_init(
         ggml_tensor * k = nullptr;
         ggml_tensor * v = nullptr;
         ggml_tensor * s = nullptr;
+        if (qnext_recurrent) {
+            s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_embd_v_s(), qnext_state_slots);
+            auto s_name = std::string{"cache_s_l"} + std::to_string(i);
+            ggml_set_name(s, s_name.c_str());
+            cache.s_l[i] = s;
+            cache.k_l.push_back(nullptr);
+            if (needs_v_cache && !is_dsv4_k_only && model.arch != LLM_ARCH_OPENPANGU) {
+                cache.v_l.push_back(nullptr);
+            }
+            LLAMA_LOG_DEBUG("=== Created recurrent cache %s as %ld x %ld x %ld x %ld\n", s->name, s->ne[0], s->ne[1], s->ne[2], s->ne[3]);
+            if ((split_cache || replicate_mla) && model.layers[i].ssm_out->extra) {
+                auto split_ssm_out = (const ggml_split_tensor_t *)model.layers[i].ssm_out->extra;
+                GGML_ASSERT(split_ssm_out);
+                int num_v_heads = hparams.ssm_dt_rank;
+                int head_v_dim  = hparams.ssm_d_inner / num_v_heads;
+                int n_device = split_ssm_out->n_device;
+                auto & split_s_l = cache.split_s_l.emplace_back();
+                split_s_l.tensor_splits.resize(n_device, nullptr);
+                for (int is = 0; is < n_device; ++is) {
+                    auto split = split_ssm_out->splits[is];
+                    if (!split) continue;
+                    GGML_ASSERT(split->ne[0] % head_v_dim == 0);
+                    int nv = split->ne[0] / head_v_dim;
+                    auto size = hparams.n_embd_v_s_id(nv);
+                    split_s_l.tensor_splits[is] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, size, qnext_state_slots);
+                    auto split_name = s_name + '.' + std::to_string(is);
+                    ggml_set_name(split_s_l.tensor_splits[is], split_name.c_str());
+                    mem_split[is] += ggml_nbytes(split_s_l.tensor_splits[is]);
+                }
+                split_s_l.ggml.n_device  = n_device;
+                split_s_l.ggml.split_dim = 0;
+                split_s_l.ggml.splits    = split_s_l.tensor_splits.data();
+                cache.s_l[i]->extra = (void *)&split_s_l.ggml;
+            }
+            continue;
+        }
         if (is_mla_attn && cparams.mla_attn) {
             // DeepSeek MLA
             const uint32_t n_embd_head_qk_rope = hparams.n_rot;
@@ -1432,40 +1471,6 @@ static bool llama_kv_cache_init(
                 }
                 continue;
             }
-            if (qnext_recurrent) {
-                s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_embd_v_s(), qnext_state_slots);
-                auto s_name = std::string{"cache_s_l"} + std::to_string(i);
-                ggml_set_name(s, s_name.c_str());
-                cache.s_l[i] = s;
-                cache.k_l.push_back(nullptr);
-                cache.v_l.push_back(nullptr);
-                LLAMA_LOG_DEBUG("=== Created recurrent cache %s as %ld x %ld x %ld x %ld\n", s->name, s->ne[0], s->ne[1], s->ne[2], s->ne[3]);
-                if (split_cache && model.layers[i].ssm_out->extra) {
-                    auto split_ssm_out = (const ggml_split_tensor_t *)model.layers[i].ssm_out->extra;
-                    GGML_ASSERT(split_ssm_out);
-                    int num_v_heads = hparams.ssm_dt_rank;
-                    int head_v_dim  = hparams.ssm_d_inner / num_v_heads;
-                    int n_device = split_ssm_out->n_device;
-                    auto & split_s_l = cache.split_s_l.emplace_back();
-                    split_s_l.tensor_splits.resize(n_device, nullptr);
-                    for (int is = 0; is < n_device; ++is) {
-                        auto split = split_ssm_out->splits[is];
-                        if (!split) continue;
-                        GGML_ASSERT(split->ne[0] % head_v_dim == 0);
-                        int nv = split->ne[0] / head_v_dim;
-                        auto size = hparams.n_embd_v_s_id(nv);
-                        split_s_l.tensor_splits[is] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, size, qnext_state_slots);
-                        auto split_name = s_name + '.' + std::to_string(is);
-                        ggml_set_name(split_s_l.tensor_splits[is], split_name.c_str());
-                        mem_split[is] += ggml_nbytes(split_s_l.tensor_splits[is]);
-                    }
-                    split_s_l.ggml.n_device  = n_device;
-                    split_s_l.ggml.split_dim = 0;
-                    split_s_l.ggml.splits    = split_s_l.tensor_splits.data();
-                    cache.s_l[i]->extra = (void *)&split_s_l.ggml;
-                }
-                continue;
-            }
             bool split_cache_i = split_cache;
             auto K = model.layers[i].wk;
             auto V = model.layers[i].wv;
@@ -1583,7 +1588,7 @@ static bool llama_kv_cache_init(
             }
         }
     }
-    if (is_mla_attn && cparams.mla_attn && n_mla < n_kv_active_layers && n_mla > 0) {
+    if (is_mla_attn && cparams.mla_attn && !llm_arch_is_hybrid(model.arch) && n_mla < n_kv_active_layers && n_mla > 0) {
         LLAMA_LOG_ERROR("%s: unexpected situation with %d out of %d active KV layers having MLA enabled\n", __func__, n_mla, n_kv_active_layers);
         LLAMA_LOG_ERROR("%s: bailing out\n", __func__);
         GGML_ABORT("fatal error");
@@ -2739,7 +2744,7 @@ static void llm_load_print_meta(llama_model_loader & ml, llama_model & model) {
         LLAMA_LOG_INFO("%s: f_attention_scale = %f\n", __func__, hparams.f_attention_scale);
     }
 
-    if (model.arch == LLM_ARCH_BAILINGMOE2) {
+    if (model.arch == LLM_ARCH_BAILINGMOE2 || model.arch == LLM_ARCH_BAILINGMOE3) {
         LLAMA_LOG_INFO("%s: n_layer_dense_lead   = %d\n",     __func__, hparams.n_layer_dense_lead);
         LLAMA_LOG_INFO("%s: n_ff_exp             = %d\n",     __func__, hparams.n_ff_exp);
         LLAMA_LOG_INFO("%s: n_ff_shexp           = %d\n",     __func__, hparams.n_ff_shexp);
@@ -8783,6 +8788,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) {
         case LLM_ARCH_MISTRAL3:
         case LLM_ARCH_GLM_DSA:
         case LLM_ARCH_MISTRAL4:
+        case LLM_ARCH_BAILINGMOE3:
         case LLM_ARCH_DFLASH:
         case LLM_ARCH_MUSE_GLIMMER:
             return LLAMA_ROPE_TYPE_NORM;