From 87644e36bcc1580f424d5b86328e36cf6464ae96 Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Thu, 13 Aug 2026 01:02:50 -0500 Subject: [PATCH] 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 | 1 + ggml/src/ggml-cuda.cu | 6 +- ggml/src/ggml-cuda/kda.cu | 268 +++++++++++++++++++++++ ggml/src/ggml-cuda/kda.cuh | 3 + ggml/src/ggml.c | 149 ++++++++++++- ggml/src/iqk/iqk_cpu_ops.h | 9 +- ggml/src/iqk/iqk_kda.cpp | 363 +++++++++++++++++++++++++++++++ src/CMakeLists.txt | 3 + src/graphs/build_bailingmoe3.cpp | 99 +++++++++ src/graphs/build_deepseek2.cpp | 28 ++- src/llama-arch.cpp | 5 + src/llama-arch.h | 11 + src/llama-build-context.cpp | 59 +++-- src/llama-build-context.h | 5 + src/llama-delta-net.cpp | 6 +- src/llama-delta-net.h | 7 + src/llama-hparams.cpp | 67 +++++- src/llama-hparams.h | 4 + src/llama-kda.cpp | 250 +++++++++++++++++++++ src/llama-kda.h | 3 + src/llama-load-tensors.cpp | 161 +++++++++++++- src/llama-model.cpp | 51 ++++- src/llama-model.h | 19 +- src/llama-quantize.cpp | 2 +- src/llama.cpp | 78 ++++--- 25 files changed, 1576 insertions(+), 81 deletions(-) create mode 100644 ggml/src/ggml-cuda/kda.cu create mode 100644 ggml/src/ggml-cuda/kda.cuh create mode 100644 ggml/src/iqk/iqk_kda.cpp create mode 100644 src/graphs/build_bailingmoe3.cpp create mode 100644 src/llama-kda.cpp create mode 100644 src/llama-kda.h diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 32ed985d..cc42d8d0 100644 --- a/ggml/src/CMakeLists.txt +++ b/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 diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index b3309d2f..3a6b6efd 100644 --- a/ggml/src/ggml-cuda.cu +++ b/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: diff --git a/ggml/src/ggml-cuda/kda.cu b/ggml/src/ggml-cuda/kda.cu new file mode 100644 index 00000000..cbb21615 --- /dev/null +++ b/ggml/src/ggml-cuda/kda.cu @@ -0,0 +1,268 @@ +#include "common.cuh" +#include "kda.cuh" + +#include + +__device__ __forceinline__ float sigmoid_f(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +template +__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 +__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(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><<>>( + 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><<>>( + 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><<>>( + 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><<>>( + 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()); +} diff --git a/ggml/src/ggml-cuda/kda.cuh b/ggml/src/ggml-cuda/kda.cuh new file mode 100644 index 00000000..bb10d25c --- /dev/null +++ b/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); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 17eb02c2..4901f37a 100644 --- a/ggml/src/ggml.c +++ b/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]; diff --git a/ggml/src/iqk/iqk_cpu_ops.h b/ggml/src/iqk/iqk_cpu_ops.h index 5847dc22..4e499481 100644 --- a/ggml/src/iqk/iqk_cpu_ops.h +++ b/ggml/src/iqk/iqk_cpu_ops.h @@ -7,6 +7,7 @@ #pragma once #include #include +#include #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 - diff --git a/ggml/src/iqk/iqk_kda.cpp b/ggml/src/iqk/iqk_kda.cpp new file mode 100644 index 00000000..d96cd3de --- /dev/null +++ b/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 +#include + +namespace { + +struct kda_layout { + size_t v[3]; + size_t g[4]; + size_t beta[4]; +}; + +#ifdef __ARM_NEON +template +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 +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(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 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 388ccb2c..6fd69a82 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/graphs/build_bailingmoe3.cpp b/src/graphs/build_bailingmoe3.cpp new file mode 100644 index 00000000..f1caf0e4 --- /dev/null +++ b/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; +} diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index 13a54347..981f9229 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/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); + } } } diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index c5078c7f..dd00a254 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -73,6 +73,7 @@ static const std::map 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_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; diff --git a/src/llama-arch.h b/src/llama-arch.h index 5fa1cb9c..7b47e1f5 100644 --- a/src/llama-arch.h +++ b/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, diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index bd52b863..5578119b 100644 --- a/src/llama-build-context.cpp +++ b/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(); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 556f7685..620beacd 100644 --- a/src/llama-build-context.h +++ b/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); diff --git a/src/llama-delta-net.cpp b/src/llama-delta-net.cpp index 7a622de9..aad31c5f 100644 --- a/src/llama-delta-net.cpp +++ b/src/llama-delta-net.cpp @@ -95,7 +95,9 @@ std::pair 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 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); diff --git a/src/llama-delta-net.h b/src/llama-delta-net.h index 89c4b130..36722d8f 100644 --- a/src/llama-delta-net.h +++ b/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; }; diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 68ac76d9..9d8ec3ae 100644 --- a/src/llama-hparams.cpp +++ b/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(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); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index cc1d43db..b68a7384 100644 --- a/src/llama-hparams.h +++ b/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 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; diff --git a/src/llama-kda.cpp b/src/llama-kda.cpp new file mode 100644 index 00000000..744666aa --- /dev/null +++ b/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 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 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 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; +} diff --git a/src/llama-kda.h b/src/llama-kda.h new file mode 100644 index 00000000..e2ed70d8 --- /dev/null +++ b/src/llama-kda.h @@ -0,0 +1,3 @@ +#pragma once + +#include "llama-delta-net.h" diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 9d0d8eda..47300df0 100644 --- a/src/llama-load-tensors.cpp +++ b/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 & cur_splits, std::vector & 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); } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index aaf088db..cbf10eb3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1663,6 +1663,55 @@ static const std::map> 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; diff --git a/src/llama-model.h b/src/llama-model.h index 891a3772..4a319590 100644 --- a/src/llama-model.h +++ b/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 diff --git a/src/llama-quantize.cpp b/src/llama-quantize.cpp index 72fd197b..7ebc37d6 100644 --- a/src/llama-quantize.cpp +++ b/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; diff --git a/src/llama.cpp b/src/llama.cpp index 9aa5f952..0fc6a9cb 100644 --- a/src/llama.cpp +++ b/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;