public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

talk-llama : sync llama.cpp

This commit is contained in:
Georgi Gerganov 2026-08-18 11:32:12 +03:00
parent 6b014cfbf6
commit 4834a2327d
22 changed files with 2017 additions and 112 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
examples/talk-llama/llama-adapter.cpp
View File

@ -396,8 +396,11 @@ static void llama_adapter_lora_init_impl(llama_model & model, const char * path_
llama_file gguf_file(path_lora, "rb");
std::vector<uint8_t> read_buf;
auto set_tensor = [&](ggml_tensor * orig, ggml_tensor * dev) {
size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name));
size_t size = ggml_nbytes(orig);
const size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name));
const size_t size = ggml_nbytes(orig);
if (offs + size < offs || offs + size > gguf_file.size()) {
throw std::runtime_error(format("LoRA tensor '%s' data is not within the file bounds, file is corrupted or incomplete", orig->name));
}
read_buf.resize(size);
gguf_file.seek(offs, SEEK_SET);
gguf_file.read_raw(read_buf.data(), size);

31
examples/talk-llama/llama-arch.cpp
View File

@ -107,6 +107,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_PLM, "plm" },
{ LLM_ARCH_BAILINGMOE, "bailingmoe" },
{ LLM_ARCH_BAILINGMOE2, "bailingmoe2" },
{ LLM_ARCH_BAILINGMOE3, "bailingmoe3" },
{ LLM_ARCH_DOTS1, "dots1" },
{ LLM_ARCH_ARCEE, "arcee" },
{ LLM_ARCH_AFMOE, "afmoe" },
@ -128,6 +129,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_SEED_OSS, "seed_oss" },
{ LLM_ARCH_GROVEMOE, "grovemoe" },
{ LLM_ARCH_APERTUS, "apertus" },
{ LLM_ARCH_MINIMAX_01, "minimax-01" },
{ LLM_ARCH_MINIMAX_M2, "minimax-m2" },
{ LLM_ARCH_MINIMAX_M3, "minimax-m3" },
{ LLM_ARCH_COGVLM, "cogvlm" },
@ -143,6 +145,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_LLAMA_EMBED, "llama-embed" },
{ LLM_ARCH_MAINCODER, "maincoder" },
{ LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
{ LLM_ARCH_KIMI_K3, "kimi-k3" },
{ LLM_ARCH_TALKIE, "talkie" },
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
@ -186,6 +189,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_FEATURES_LENGTH, "%s.features_length" },
{ LLM_KV_BLOCK_COUNT, "%s.block_count" },
{ LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" },
{ LLM_KV_ATTN_RES_BLOCK_SIZE, "%s.attn_res.block_size" },
{ LLM_KV_ACTIVATION_SITU_BETA, "%s.activation.situ_beta" },
{ LLM_KV_ACTIVATION_SITU_LINEAR_BETA, "%s.activation.situ_linear_beta" },
{ LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" },
{ LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" },
{ LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" },
@ -201,6 +207,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_EXPERT_GROUP_USED_COUNT, "%s.expert_group_used_count" },
{ LLM_KV_EXPERT_WEIGHTS_SCALE, "%s.expert_weights_scale" },
{ LLM_KV_EXPERT_WEIGHTS_NORM, "%s.expert_weights_norm" },
{ LLM_KV_EXPERT_LATENT_LENGTH, "%s.expert_latent_length" },
{ LLM_KV_EXPERT_GATING_FUNC, "%s.expert_gating_func" },
{ LLM_KV_EXPERT_GROUP_SCALE, "%s.expert_group_scale" },
{ LLM_KV_EXPERTS_PER_GROUP, "%s.experts_per_group" },
@ -311,7 +318,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_SSM_GROUP_COUNT, "%s.ssm.group_count" },
{ LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" },
{ LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" },
{ 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_WKV_HEAD_SIZE, "%s.wkv.head_size" },
@ -462,6 +471,13 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_SSM_F_B, "blk.%d.ssm_f_b" },
{ LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" },
{ LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" },
{ LLM_TENSOR_SSM_G, "blk.%d.ssm_g" },
{ LLM_TENSOR_ATTN_RES_SCORE, "blk.%d.attn_res_score" },
{ LLM_TENSOR_FFN_RES_SCORE, "blk.%d.ffn_res_score" },
{ LLM_TENSOR_OUTPUT_RES_SCORE, "output_res_score" },
{ LLM_TENSOR_FFN_ROUTED_DOWN, "blk.%d.ffn_routed_down" },
{ LLM_TENSOR_FFN_ROUTED_UP, "blk.%d.ffn_routed_up" },
{ LLM_TENSOR_FFN_ROUTED_NORM, "blk.%d.ffn_routed_norm" },
{ LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" },
{ LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" },
{ LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" },
@ -755,6 +771,13 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_SSM_F_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_SSM_G, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_FFN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_OUTPUT_RES_SCORE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}},
{LLM_TENSOR_FFN_ROUTED_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_FFN_ROUTED_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_FFN_ROUTED_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
@ -975,9 +998,12 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_QWEN3NEXT:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_DEEPSEEK4:
case LLM_ARCH_MINIMAX_01:
return true;
default:
return false;
@ -1033,10 +1059,13 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_GRANITE_HYBRID:
case LLM_ARCH_LFM2:
case LLM_ARCH_LFM2MOE:
case LLM_ARCH_MINIMAX_01:
case LLM_ARCH_MINIMAX_M2:
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_QWEN3TTS:
return false;
default:

16
examples/talk-llama/llama-arch.h
View File

@ -112,6 +112,7 @@ enum llm_arch {
LLM_ARCH_PLM,
LLM_ARCH_BAILINGMOE,
LLM_ARCH_BAILINGMOE2,
LLM_ARCH_BAILINGMOE3,
LLM_ARCH_DOTS1,
LLM_ARCH_ARCEE,
LLM_ARCH_AFMOE,
@ -145,6 +146,7 @@ enum llm_arch {
LLM_ARCH_LLAMA_EMBED,
LLM_ARCH_MAINCODER,
LLM_ARCH_KIMI_LINEAR,
LLM_ARCH_KIMI_K3,
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
@ -153,6 +155,7 @@ enum llm_arch {
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
LLM_ARCH_POCKETTTS,
LLM_ARCH_MINIMAX_01,
LLM_ARCH_UNKNOWN,
};
@ -191,6 +194,9 @@ enum llm_kv {
LLM_KV_FEATURES_LENGTH,
LLM_KV_BLOCK_COUNT,
LLM_KV_LEADING_DENSE_BLOCK_COUNT,
LLM_KV_ATTN_RES_BLOCK_SIZE,
LLM_KV_ACTIVATION_SITU_BETA,
LLM_KV_ACTIVATION_SITU_LINEAR_BETA,
LLM_KV_FEED_FORWARD_LENGTH,
LLM_KV_EXPERT_FEED_FORWARD_LENGTH,
LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH,
@ -206,6 +212,7 @@ enum llm_kv {
LLM_KV_EXPERT_GROUP_USED_COUNT,
LLM_KV_EXPERT_WEIGHTS_SCALE,
LLM_KV_EXPERT_WEIGHTS_NORM,
LLM_KV_EXPERT_LATENT_LENGTH,
LLM_KV_EXPERT_GATING_FUNC,
LLM_KV_EXPERT_GROUP_SCALE,
LLM_KV_EXPERTS_PER_GROUP,
@ -317,6 +324,8 @@ enum llm_kv {
LLM_KV_SSM_DT_B_C_RMS,
LLM_KV_KDA_HEAD_DIM,
LLM_KV_KDA_SAFE_GATE,
LLM_KV_KDA_GATE_LOWER_BOUND,
LLM_KV_WKV_HEAD_SIZE,
@ -491,6 +500,13 @@ enum llm_tensor {
LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient and qwen3.5
LLM_TENSOR_SSM_G_A, // kimi: output gate projection A
LLM_TENSOR_SSM_G_B, // kimi: output gate projection B
LLM_TENSOR_SSM_G, // kimi-k3: full-rank KDA gate
LLM_TENSOR_ATTN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-attn)
LLM_TENSOR_FFN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-ffn)
LLM_TENSOR_OUTPUT_RES_SCORE, // kimi-k3: fused res_norm*res_proj (final)
LLM_TENSOR_FFN_ROUTED_DOWN, // kimi-k3: latent MoE down
LLM_TENSOR_FFN_ROUTED_UP, // kimi-k3: latent MoE up
LLM_TENSOR_FFN_ROUTED_NORM, // kimi-k3: latent MoE norm
LLM_TENSOR_TIME_MIX_W0,
LLM_TENSOR_TIME_MIX_W1,
LLM_TENSOR_TIME_MIX_W2,

7
examples/talk-llama/llama-context.cpp
View File

@ -2293,13 +2293,18 @@ void llama_context::output_reorder() {
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
uint32_t res;
if (model.arch == LLM_ARCH_QWEN3NEXT ||
if (model.arch == LLM_ARCH_KIMI_K3) {
// the n_tokens*40 budget below is exhausted at ubatch 3840
res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors());
} else if (model.arch == LLM_ARCH_QWEN3NEXT ||
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_BAILINGMOE3 ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
model.arch == LLM_ARCH_DEEPSEEK4 ||
(model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
model.arch == LLM_ARCH_NANBEIGE ||
model.arch == LLM_ARCH_MINIMAX_01 ||
model.arch == LLM_ARCH_MINIMAX_M3) {
res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
} else {

17
examples/talk-llama/llama-graph.cpp
View File

@ -1835,6 +1835,8 @@ ggml_tensor * llm_graph_context::build_ffn(
cur = ggml_reglu(ctx0, cur);
cb(cur, "ffn_reglu", il);
} break;
case LLM_FFN_SITU:
GGML_ABORT("not yet supported");
default:
GGML_ABORT("fatal error");
}
@ -2174,6 +2176,21 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
cur = ggml_silu(ctx0, cur);
cb(cur, "ffn_moe_silu", il);
} break;
case LLM_FFN_SITU:
{
// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb)
GGML_ASSERT(has_gate);
const float beta = hparams.situ_beta;
const float lb = hparams.situ_linear_beta;
ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta);
act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur));
if (lb > 0.0f) {
up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb);
}
cur = ggml_mul(ctx0, act, up);
cb(cur, "ffn_moe_situ", il);
} break;
case LLM_FFN_GELU:
if (has_gate) {
cur = ggml_geglu_split(ctx0, cur, up);

1
examples/talk-llama/llama-graph.h
View File

@ -59,6 +59,7 @@ enum llm_ffn_op_type : int {
LLM_FFN_GEGLU,
LLM_FFN_REGLU,
LLM_FFN_SWIGLU_OAI_MOE,
LLM_FFN_SITU, // kimi-k3
};
enum llm_ffn_gate_type {

7
examples/talk-llama/llama-hparams.cpp
View File

@ -217,6 +217,13 @@ uint32_t llama_hparams::n_embd_s() const {
return n_embd_head_kda * n_embd_head_kda * n_head(); // 128 * 128 * 32 = 524288
}
if (n_embd_head_la != 0) {
// for MiniMax-Text-01 linear attention layers
// Full recurrent state: head_dim * head_dim * n_head
// tensor shape for linear attention: [head_dim, head_dim, n_head]
return n_embd_head_la * n_embd_head_la * n_head(); // 128 * 128 * 64 = 1048576
}
// corresponds to Mamba's ssm_states size
return ssm_d_state * ssm_d_inner;
}

14
examples/talk-llama/llama-hparams.h
View File

@ -4,10 +4,11 @@
#include <array>
#include <cassert>
#include <cmath>
// bump if necessary
#define LLAMA_MAX_LAYERS 512
#define LLAMA_MAX_EXPERTS 512 // Qwen3 Next
#define LLAMA_MAX_EXPERTS 1024 // Kimi K3
enum llama_expert_gating_func_type {
LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0,
@ -164,8 +165,19 @@ struct llama_hparams {
uint32_t ssm_dt_rank = 0;
uint32_t ssm_n_group = 0;
// for MiniMax-Text-01 linear attention
uint32_t n_embd_head_la = 0;
// for Kimi Linear KDA
uint32_t n_embd_head_kda = 0;
bool kda_safe_gate = false;
// kimi-k3
uint32_t n_expert_latent = 0; // routed_expert_hidden_size (0 = experts run at n_embd)
uint32_t attn_res_block_size = 0; // 0 = no cross-layer attention residuals
float kda_gate_lower_bound = -INFINITY;
float situ_beta = 1.0f;
float situ_linear_beta = 0.0f; // 0 = no linear-beta transform on the up branch
bool ssm_dt_b_c_rms = false;

26
examples/talk-llama/llama-model-loader.cpp
View File

@ -316,15 +316,19 @@ namespace GGUFMeta {
struct GGUFMeta::ArrayInfo arr_info =
GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
bool type_ok = false;
switch (arr_info.gt) {
case GGUF_TYPE_UINT32:
case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same<T, int32_t>::value) ||
(std::is_same<T, uint32_t>::value)); break;
case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T, float>::value)); break;
case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same<T, std::string>::value)); break;
case GGUF_TYPE_INT32: type_ok = (std::is_same<T, int32_t>::value) ||
(std::is_same<T, uint32_t>::value); break;
case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T, float>::value); break;
case GGUF_TYPE_STRING: type_ok = (std::is_same<T, std::string>::value); break;
default:
throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
}
if (!type_ok) {
throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
}
if constexpr (std::is_same<T, std::string>::value) {
const size_t n_items = gguf_get_arr_n(ctx, kid);
@ -357,16 +361,20 @@ namespace GGUFMeta {
struct GGUFMeta::ArrayInfo arr_info =
GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
bool type_ok = false;
switch (arr_info.gt) {
case GGUF_TYPE_BOOL:
case GGUF_TYPE_UINT32:
case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same<T, int32_t>::value) ||
(std::is_same<T, uint32_t>::value)); break;
case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T, float>::value)); break;
case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same<T, std::string>::value)); break;
case GGUF_TYPE_INT32: type_ok = (std::is_same<T, int32_t>::value) ||
(std::is_same<T, uint32_t>::value); break;
case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T, float>::value); break;
case GGUF_TYPE_STRING: type_ok = (std::is_same<T, std::string>::value); break;
default:
throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
}
if (!type_ok) {
throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
}
if (arr_info.length > N_MAX) {
throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX));
@ -1178,7 +1186,7 @@ struct ggml_tensor * llama_model_loader::create_tensor(
if (use_mmap) {
static std::once_flag once;
std::call_once(once, [] {
LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance\n");
LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --load-mode none for better performance\n");
});
}
} else {

15
examples/talk-llama/llama-model-saver.cpp
View File

@ -121,6 +121,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c
}
// instantiate for external usage:
template void llama_model_saver::add_kv<std::vector<uint32_t>>(const enum llm_kv, const std::vector<uint32_t> &, const bool);
template void llama_model_saver::add_kv<std::vector<float>>(const enum llm_kv, const std::vector<float> &, const bool);
void llama_model_saver::add_kv(const enum llm_kv key, const std::vector<std::string> & value) {
std::vector<const char *> tmp(value.size());
@ -213,10 +214,13 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead);
add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true);
add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
add_kv(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent);
add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp);
add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp);
add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector<float>(
hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.begin() + hparams.n_layer_all));
add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector<float>(
hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.begin() + hparams.n_layer_all));
add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res);
// add_kv(LLM_KV_TENSOR_DATA_LAYOUT, ???);
add_kv(LLM_KV_EXPERT_COUNT, hparams.n_expert);
@ -319,6 +323,8 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms);
add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
add_kv(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate);
add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound);
add_kv(LLM_KV_WKV_HEAD_SIZE, hparams.wkv_head_size);
@ -376,6 +382,10 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_XIELU_BETA, hparams.xielu_beta);
add_kv(LLM_KV_XIELU_EPS, hparams.xielu_eps);
add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size);
add_kv(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta);
add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta);
// deprecated
// add_kv(LLM_KV_TOKENIZER_PREFIX_ID, ???);
// add_kv(LLM_KV_TOKENIZER_SUFFIX_ID, ???);
@ -403,6 +413,7 @@ void llama_model_saver::add_tensors_from_model() {
add_tensor(model->output_norm_enc);
add_tensor(model->output_s);
add_tensor(model->output_in_s);
add_tensor(model->output_res_score);
add_tensor(model->cls);
add_tensor(model->cls_b);
add_tensor(model->cls_out);

23
examples/talk-llama/llama-model.cpp
View File

@ -256,6 +256,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_bailingmoe(params);
case LLM_ARCH_BAILINGMOE2:
return new llama_model_bailingmoe2(params);
case LLM_ARCH_BAILINGMOE3:
return new llama_model_bailingmoe3(params);
case LLM_ARCH_SEED_OSS:
return new llama_model_seed_oss(params);
case LLM_ARCH_DOTS1:
@ -296,6 +298,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_grovemoe(params);
case LLM_ARCH_APERTUS:
return new llama_model_apertus(params);
case LLM_ARCH_MINIMAX_01:
return new llama_model_minimax_01(params);
case LLM_ARCH_MINIMAX_M2:
return new llama_model_minimax_m2(params);
case LLM_ARCH_MINIMAX_M3:
@ -320,6 +324,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_mimo2(params);
case LLM_ARCH_KIMI_LINEAR:
return new llama_model_kimi_linear(params);
case LLM_ARCH_KIMI_K3:
return new llama_model_kimi_k3(params);
case LLM_ARCH_STEP35:
return new llama_model_step35(params);
default:
@ -798,6 +804,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_290B: return "290B";
case LLM_TYPE_314B: return "314B";
case LLM_TYPE_405B: return "405B";
case LLM_TYPE_456B: return "456B";
case LLM_TYPE_671B: return "671B";
case LLM_TYPE_SMALL: return "0.1B";
case LLM_TYPE_MEDIUM: return "0.4B";
@ -816,6 +823,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_A13B: return "A13B";
case LLM_TYPE_7B_A1B: return "7B.A1B";
case LLM_TYPE_8B_A1B: return "8B.A1B";
case LLM_TYPE_7_9B_A1_3B: return "7.9B.A1.3B";
case LLM_TYPE_12B_A2_5B: return "12B.A2.5B";
case LLM_TYPE_16B_A1B: return "16B.A1B";
case LLM_TYPE_21B_A3B: return "21B.A3B";
@ -832,6 +840,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_118B_A8B: return "118B.A8B";
case LLM_TYPE_120B_A12B: return "120B.A12B";
case LLM_TYPE_122B_A10B: return "122B.A10B";
case LLM_TYPE_124B_A5_1B: return "124B.A5.1B";
case LLM_TYPE_196B_A11B: return "196B.A11B";
case LLM_TYPE_230B_A10B: return "230B.A10B";
case LLM_TYPE_428B_A23B: return "428B.A23B";
@ -842,6 +851,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_397B_A17B: return "397B.A17B";
case LLM_TYPE_685B_A37B: return "685B.A37B";
case LLM_TYPE_744B_A40B: return "744B.A40B";
case LLM_TYPE_2_8T_A50B: return "2.8T.A50B";
case LLM_TYPE_E2B: return "E2B";
case LLM_TYPE_E4B: return "E4B";
default: return "?B";
@ -1954,7 +1964,7 @@ void llama_model::print_info() const {
LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm);
}
if (arch == LLM_ARCH_BAILINGMOE2) {
if (arch == LLM_ARCH_BAILINGMOE2 || 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);
@ -2249,11 +2259,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
// checks
default:
{
// The MTP head is dense-attention only on hybrid Qwen3-Next/3.5/3.6, so use a plain
// attention KV cache for the MTP context instead of the hybrid wrapper.
// Dense MTP heads use a plain attention KV cache instead of the hybrid wrapper.
const bool mtp_on_hybrid_qwen =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ||
arch == LLM_ARCH_BAILINGMOE3);
const bool mtp_on_hybrid_nemotron =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE;
@ -2283,7 +2293,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
filter_recr = [&](uint32_t il) {
return hparams.is_recr(il) && hparams.n_ff(il) == 0;
};
} else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
} else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) {
filter_attn = [&](uint32_t il) {
return il < hparams.n_layer() && !hparams.is_recr(il);
};
@ -2599,6 +2609,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
return LLAMA_ROPE_TYPE_NONE;
// use what we call a normal RoPE, operating on pairs of consecutive head values
@ -2630,6 +2641,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_GRANITE_SWITCH:
case LLM_ARCH_CHAMELEON:
case LLM_ARCH_BAILINGMOE:
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_NEO_BERT:
case LLM_ARCH_SMOLLM3:
case LLM_ARCH_ARCEE:
@ -2704,6 +2716,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_SEED_OSS:
case LLM_ARCH_GROVEMOE:
case LLM_ARCH_APERTUS:
case LLM_ARCH_MINIMAX_01:
case LLM_ARCH_MINIMAX_M2:
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_COGVLM:

14
examples/talk-llama/llama-model.h
View File

@ -99,6 +99,7 @@ enum llm_type {
LLM_TYPE_290B,
LLM_TYPE_314B,
LLM_TYPE_405B,
LLM_TYPE_456B,
LLM_TYPE_671B,
LLM_TYPE_SMALL,
LLM_TYPE_MEDIUM,
@ -117,6 +118,7 @@ enum llm_type {
LLM_TYPE_A13B,
LLM_TYPE_7B_A1B,
LLM_TYPE_8B_A1B, // lfm2moe
LLM_TYPE_7_9B_A1_3B, // Ling-3.0-tiny
LLM_TYPE_12B_A2_5B,
LLM_TYPE_16B_A1B,
LLM_TYPE_21B_A3B, // Ernie MoE small
@ -133,6 +135,7 @@ enum llm_type {
LLM_TYPE_118B_A8B, // Laguna-S-2
LLM_TYPE_120B_A12B, // Nemotron 3 Super
LLM_TYPE_122B_A10B, // Qwen3.5
LLM_TYPE_124B_A5_1B, // Ling-3.0-flash
LLM_TYPE_196B_A11B, // Step3.5-Flash
LLM_TYPE_230B_A10B, // Minimax M2
LLM_TYPE_428B_A23B, // Minimax M3
@ -143,6 +146,7 @@ enum llm_type {
LLM_TYPE_397B_A17B, // Qwen3.5
LLM_TYPE_685B_A37B, // DeepSeek V3.2
LLM_TYPE_744B_A40B, // GLM-5
LLM_TYPE_2_8T_A50B, // Kimi-K3
LLM_TYPE_E2B,
LLM_TYPE_E4B,
};
@ -271,6 +275,7 @@ struct llama_layer {
struct ggml_tensor * wv = nullptr;
struct ggml_tensor * wo = nullptr;
struct ggml_tensor * wqkv = nullptr;
struct ggml_tensor * wg = nullptr;
struct ggml_tensor * wq_a = nullptr;
struct ggml_tensor * wq_b = nullptr;
struct ggml_tensor * wkv_a_mqa = nullptr;
@ -528,6 +533,14 @@ struct llama_layer {
struct ggml_tensor * ssm_g_b = nullptr;
struct ggml_tensor * ssm_o_norm = nullptr;
// kimi-k3
struct ggml_tensor * ssm_g = nullptr; // full-rank KDA gate (replaces ssm_g_a/ssm_g_b)
struct ggml_tensor * attn_res_score = nullptr; // fused res_norm*res_proj, pre-attention
struct ggml_tensor * ffn_res_score = nullptr; // fused res_norm*res_proj, pre-FFN
struct ggml_tensor * ffn_routed_down = nullptr; // latent MoE: n_embd -> n_expert_latent
struct ggml_tensor * ffn_routed_up = nullptr; // latent MoE: n_expert_latent -> n_embd
struct ggml_tensor * ffn_routed_norm = nullptr;
// DSA (deepseek sparse attention)
struct ggml_tensor * indexer_k_norm = nullptr;
struct ggml_tensor * indexer_k_norm_b = nullptr;
@ -587,6 +600,7 @@ struct llama_model {
struct ggml_tensor * tok_norm_b = nullptr;
struct ggml_tensor * output_norm = nullptr;
struct ggml_tensor * output_res_score = nullptr; // kimi-k3: final cross-layer residual mix
struct ggml_tensor * output_norm_b = nullptr;
struct ggml_tensor * output = nullptr;
struct ggml_tensor * output_b = nullptr;

7
examples/talk-llama/llama-quant.cpp
View File

@ -474,7 +474,12 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type
} else if (ftype == LLAMA_FTYPE_MOSTLY_MXFP4_MOE) {
// MoE tensors -> MXFP4
// other tensors -> Q8_0
if (tensor->ne[2] > 1) {
// MLA projection tensors are also 3D, so match expert tensor roles explicitly.
const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 &&
(category == tensor_category::FFN_UP ||
category == tensor_category::FFN_GATE ||
category == tensor_category::FFN_DOWN);
if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) {
new_type = GGML_TYPE_MXFP4;
} else {
new_type = GGML_TYPE_Q8_0;

49
examples/talk-llama/llama-vocab.cpp
View File

@ -1989,6 +1989,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
// Kimi-K2 doesn't need merges, skip
LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__);
} else {
if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
}
const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
for (int i = 0; i < n_merges; i++) {
const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
@ -2028,8 +2032,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str());
if (precompiled_charsmap_keyidx != -1) {
if (gguf_get_kv_type(ctx, precompiled_charsmap_keyidx) != GGUF_TYPE_ARRAY) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
}
const gguf_type pc_type = gguf_get_arr_type(ctx, precompiled_charsmap_keyidx);
GGML_ASSERT(pc_type == GGUF_TYPE_INT8 || pc_type == GGUF_TYPE_UINT8);
if (pc_type != GGUF_TYPE_INT8 && pc_type != GGUF_TYPE_UINT8) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
}
const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
@ -2081,6 +2090,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
throw std::runtime_error("cannot find tokenizer merges in model file\n");
}
{
if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
}
const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
for (int i = 0; i < n_merges; i++) {
const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
@ -2407,21 +2420,41 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
throw std::runtime_error("cannot find tokenizer vocab in model file\n");
}
if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_LIST).c_str()));
}
const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx);
const float * scores = nullptr;
const int * iscores = nullptr;
const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str());
if (score_idx != -1) {
const gguf_type kv_type = gguf_get_kv_type(ctx, score_idx);
const gguf_type arr_type = kv_type == GGUF_TYPE_ARRAY ? gguf_get_arr_type(ctx, score_idx) : GGUF_TYPE_COUNT;
if (arr_type != GGUF_TYPE_INT32 &&
arr_type != GGUF_TYPE_FLOAT32) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str()));
}
const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx);
if (n_scores < n_tokens) {
throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n");
}
scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
if (arr_type == GGUF_TYPE_INT32) {
iscores = (const int *) gguf_get_arr_data(ctx, score_idx);
} else {
scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
}
}
const int * toktypes = nullptr;
const int toktype_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str());
if (toktype_idx != -1) {
if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str()));
}
const uint32_t n_toktypes = gguf_get_arr_n(ctx, toktype_idx);
if (n_toktypes < n_tokens) {
throw std::runtime_error("Index out of array bounds for toktypes (" + std::to_string(n_toktypes) + " < " + std::to_string(n_tokens) + ")\n");
@ -2443,7 +2476,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
auto & token_data = id_to_token[i];
token_data.text = std::move(word);
token_data.score = scores ? scores[i] : 0.0f;
if (scores) {
token_data.score = scores[i];
} else if (iscores) {
token_data.score = static_cast<float>(iscores[i]);
} else {
token_data.score = 0.0f;
}
token_data.attr = LLAMA_TOKEN_ATTR_NORMAL;
if (toktypes) { //TODO: remove, required until per token attributes are available from GGUF file
@ -2584,6 +2623,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
{
const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str());
if (suppress_idx != -1) {
if (gguf_get_kv_type(ctx, suppress_idx) != GGUF_TYPE_ARRAY ||
gguf_get_arr_type(ctx, suppress_idx) != GGUF_TYPE_INT32) {
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str()));
}
const int n = gguf_get_arr_n(ctx, suppress_idx);
const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
// drop out-of-range ids

532
examples/talk-llama/models/bailingmoe3.cpp Normal file
View File

@ -0,0 +1,532 @@
#include "models.h"
#include "llama-memory-recurrent.h"
void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false);
ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
if (!ml.get_key(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate, false)) {
hparams.kda_safe_gate = true;
}
ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound);
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, false);
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false);
ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false);
if (hparams.n_ff_shexp == 0) {
hparams.n_ff_shexp = hparams.n_ff_exp * std::max(1u, hparams.n_expert_shared);
}
GGML_ASSERT(hparams.kda_safe_gate);
GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f);
for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0;
}
switch (hparams.n_layer()) {
case 24: type = hparams.n_embd == 1536 && hparams.n_expert == 128 ? LLM_TYPE_7_9B_A1_3B : LLM_TYPE_UNKNOWN; break;
case 42: type = hparams.n_embd == 2560 && hparams.n_expert == 512 ? LLM_TYPE_124B_A5_1B : LLM_TYPE_UNKNOWN; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
if (output == nullptr) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
}
const int64_t head_dim = hparams.n_embd_head_kda;
const int64_t d_inner = head_dim * n_head;
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t kv_lora_rank = hparams.n_lora_kv;
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t qk_rope_head_dim = hparams.n_rot();
const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
const int64_t v_head_dim = hparams.n_embd_head_v_mla();
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
for (int il = 0; il < n_layer; ++il) {
auto & layer = layers[il];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, trunk_flags);
if (hparams.is_recr(il)) {
layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags);
layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags);
layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags);
create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, trunk_flags);
layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), { n_embd, d_inner }, trunk_flags);
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_head }, trunk_flags);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), { 1, n_head }, trunk_flags);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { d_inner }, trunk_flags);
layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), { n_embd, d_inner }, trunk_flags);
layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_dim }, trunk_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { d_inner, n_embd }, trunk_flags);
} else {
if (q_lora_rank > 0) {
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, trunk_flags);
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, trunk_flags);
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, trunk_flags);
} else {
layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, trunk_flags);
}
layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, trunk_flags);
layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, trunk_flags);
layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, trunk_flags);
layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, trunk_flags);
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, trunk_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, trunk_flags);
}
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, trunk_flags);
if ((uint32_t) il < hparams.n_layer_dense_lead) {
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), { n_embd, n_ff }, trunk_flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), { n_embd, n_ff }, trunk_flags);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd }, trunk_flags);
} else {
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, trunk_flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, trunk_flags);
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, trunk_flags);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, trunk_flags);
}
}
for (int il = n_layer; il < n_layer_all; ++il) {
auto & layer = layers[il];
const int flags = mtp_flags;
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
if (q_lora_rank > 0) {
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, flags);
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, flags);
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, flags);
} else {
layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, flags);
}
layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, flags);
layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, flags);
layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, flags);
layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, flags);
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, flags);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, flags);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, flags);
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, flags);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, flags);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", il), { n_embd }, flags);
}
}
std::unique_ptr<llm_graph_context> llama_model_bailingmoe3::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
static ggml_tensor * bailingmoe3_causal_conv1d(
ggml_cgraph * gf,
ggml_context * ctx0,
ggml_tensor * conv_states_all,
ggml_tensor * conv_state_all,
int64_t qkv,
ggml_tensor * x,
ggml_tensor * proj_w,
ggml_tensor * conv_w,
int64_t d_conv,
int64_t head_dim,
int64_t n_head,
int64_t n_seq_tokens,
int64_t n_seqs,
int64_t n_tokens,
int64_t cache_head) {
const int64_t d_inner = head_dim * n_head;
const int64_t conv_state_size = (d_conv - 1) * d_inner;
const int64_t total_state_size = 3 * conv_state_size;
ggml_tensor * conv_state = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
(d_conv - 1) * ggml_element_size(conv_state_all),
total_state_size * ggml_element_size(conv_state_all),
qkv * conv_state_size * ggml_element_size(conv_state_all));
ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x,
ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
(d_conv - 1) * ggml_element_size(conv_states_all),
total_state_size * ggml_element_size(conv_states_all),
(cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight);
out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens));
return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs);
}
llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_build_delta_net_base(params), model(model) {
ggml_tensor * inpL = build_inp_embd(model.tok_embd);
cb(inpL, "model.input_embed", -1);
auto * inp = build_inp_mem_hybrid_k();
auto * inp_rs = inp->get_recr();
auto * inp_attn = inp->get_attn();
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const int64_t n_head = hparams.n_head();
const int64_t head_dim = hparams.n_embd_head_kda;
const int64_t d_inner = n_head * head_dim;
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
const int64_t v_head_dim = hparams.n_embd_head_v_mla();
const int64_t qk_rope_head_dim = hparams.n_rot();
const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim;
const int64_t kv_lora_rank = hparams.n_lora_kv;
const float kq_scale = 1.0f / sqrtf((float) qk_head_dim);
GGML_ASSERT(n_seqs > 0);
GGML_ASSERT(ubatch.equal_seqs());
GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
ggml_tensor * inpSA = inpL;
ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
if (hparams.is_recr(il)) {
const auto * mctx_cur = inp_rs->mctx;
const auto cache_head = mctx_cur->get_head();
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
ggml_tensor * q = bailingmoe3_causal_conv1d(
gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
ggml_tensor * k = bailingmoe3_causal_conv1d(
gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
ggml_tensor * v = bailingmoe3_causal_conv1d(
gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head);
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
gate = ggml_add(ctx0, gate, layer.ssm_dt_b);
gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens);
ggml_tensor * a = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1);
gate = ggml_scale(ctx0, ggml_sigmoid(ctx0, ggml_mul(ctx0, gate, a)), hparams.kda_gate_lower_bound);
gate = ggml_reshape_4d(ctx0, gate, head_dim, n_head, n_seq_tokens, n_seqs);
cb(gate, "kda_gate", il);
ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs));
q = ggml_l2_norm(ctx0, q, hparams.f_norm_rms_eps);
k = ggml_l2_norm(ctx0, k, hparams.f_norm_rms_eps);
ggml_tensor * states_all = mctx_cur->get_s_l(il);
ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs);
state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs);
auto result = build_delta_net(q, k, v, gate, beta, state, il);
ggml_tensor * out = ggml_cont(ctx0, result.first);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second,
ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs,
cache_head * hparams.n_embd_s() * ggml_element_size(states_all))));
ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur);
out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens);
out = ggml_reshape_3d(ctx0, out, head_dim, n_head, n_tokens);
out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, out_gate));
cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens));
cb(cur, "kda_out", il);
} else {
ggml_tensor * attn_input = cur;
ggml_tensor * q_all;
if (layer.wq_a) {
q_all = ggml_mul_mat(ctx0, layer.wq_a, cur);
cb(q_all, "q_a", il);
q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
cb(q_all, "q_a_norm", il);
q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all);
cb(q_all, "q_b", il);
} else {
q_all = ggml_mul_mat(ctx0, layer.wq, cur);
}
ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens,
ggml_row_size(q_all->type, qk_head_dim),
ggml_row_size(q_all->type, qk_head_dim) * n_head, 0);
ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens,
ggml_row_size(q_all->type, qk_head_dim),
ggml_row_size(q_all->type, qk_head_dim) * n_head,
ggml_row_size(q_all->type, qk_nope_head_dim));
ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens,
ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0);
ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens,
ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
ggml_row_size(kv_all->type, kv_lora_rank));
q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0);
kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens);
ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0);
cur = build_attn(inp_attn, nullptr, nullptr, nullptr,
q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il);
ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input);
attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens));
cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens);
cur = ggml_mul(ctx0, cur, attn_gate);
cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens));
cb(cur, "mla_out", il);
}
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
if ((uint32_t) il < hparams.n_layer_dense_lead) {
cur = build_ffn(cur,
layer.ffn_up, nullptr, nullptr,
layer.ffn_gate, nullptr, nullptr,
layer.ffn_down, nullptr, nullptr,
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
} else {
ggml_tensor * moe = build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_SILU,
hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il);
ggml_tensor * shared = build_ffn(cur,
layer.ffn_up_shexp, nullptr, nullptr,
layer.ffn_gate_shexp, nullptr, nullptr,
layer.ffn_down_shexp, nullptr, nullptr,
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
cur = ggml_add(ctx0, moe, shared);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
inpL = cur;
}
ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = ggml_mul_mat(ctx0, model.output, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn == 1 && "BailingMoE3 MTP requires one NextN layer");
const int il = hparams.n_layer() + cparams.nextn_layer_offset;
GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
"nextn_layer_offset out of range");
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
GGML_ASSERT(layer.nextn.shared_head_norm && "MTP block missing final norm");
const int64_t n_head = hparams.n_head();
const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
const int64_t v_head_dim = hparams.n_embd_head_v_mla();
const int64_t qk_rope_head_dim = hparams.n_rot();
const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim;
const int64_t kv_lora_rank = hparams.n_lora_kv;
const float kq_scale = 1.0f / sqrtf((float) qk_head_dim);
auto inp = std::make_unique<llm_graph_input_embd>(hparams.n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->embd);
ggml_set_name(inp->embd, "mtp_h_input");
ggml_tensor * tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens);
ggml_tensor * h_norm = build_norm(inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
ggml_tensor * cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0));
cb(cur, "mtp_eh_proj", il);
res->add_input(std::move(inp));
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
auto * inp_attn = build_attn_inp_k();
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
ggml_tensor * attn_input = cur;
ggml_tensor * q_all;
if (layer.wq_a) {
q_all = ggml_mul_mat(ctx0, layer.wq_a, cur);
cb(q_all, "q_a", il);
q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
cb(q_all, "q_a_norm", il);
q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all);
cb(q_all, "q_b", il);
} else {
q_all = ggml_mul_mat(ctx0, layer.wq, cur);
}
ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens,
ggml_row_size(q_all->type, qk_head_dim),
ggml_row_size(q_all->type, qk_head_dim) * n_head, 0);
ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens,
ggml_row_size(q_all->type, qk_head_dim),
ggml_row_size(q_all->type, qk_head_dim) * n_head,
ggml_row_size(q_all->type, qk_nope_head_dim));
ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens,
ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0);
ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens,
ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
ggml_row_size(kv_all->type, kv_lora_rank));
q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0);
kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens);
ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0);
cur = build_attn(inp_attn, nullptr, nullptr, nullptr,
q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il);
ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input);
attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens));
cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens);
cur = ggml_mul(ctx0, cur, attn_gate);
cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens));
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
ggml_tensor * moe = build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_SILU,
hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il);
ggml_tensor * shared = build_ffn(cur,
layer.ffn_up_shexp, nullptr, nullptr,
layer.ffn_gate_shexp, nullptr, nullptr,
layer.ffn_down_shexp, nullptr, nullptr,
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
cur = ggml_add(ctx0, moe, shared);
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cur = ggml_mul_mat(ctx0, model.output, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}

51
examples/talk-llama/models/deepseek32.cpp
View File

@ -180,10 +180,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
const int64_t n_indexer_head = hparams.indexer_n_head;
const int64_t n_embd_indexer_head = hparams.indexer_head_size;
const int64_t n_embd_indexer_head_rope = hparams.n_rot();
const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope;
const uint32_t n_indexer_top_k = hparams.indexer_top_k;
// the indexer head layous is [rope | nope]
GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head);
const uint32_t kv_lora_rank = hparams.n_lora_kv;
// We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
@ -233,28 +234,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr);
cb(indexer_q, "indexer_q", il);
// split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens}
ggml_tensor * indexer_q_pe =
ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens,
ggml_row_size(indexer_q->type, n_embd_indexer_head),
ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0);
cb(indexer_q_pe, "indexer_q_pe", il);
// and {n_embd_indexer_head_nope, n_indexer_head, n_tokens}
ggml_tensor * indexer_q_nope =
ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens,
ggml_row_size(indexer_q->type, n_embd_indexer_head),
ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head,
ggml_row_size(indexer_q->type, n_embd_indexer_head_nope));
cb(indexer_q_nope, "indexer_q_nope", il);
indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot,
// {n_embd_indexer_head, n_indexer_head, n_tokens}
indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens);
indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot,
LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(indexer_q_pe, "indexer_q_pe", il);
// {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens}
indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0);
cb(indexer_q, "indexer_q", il);
ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur);
@ -263,28 +247,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il);
cb(indexer_k, "indexer_k", il);
// split into {n_embd_indexer_head_rope, 1, n_tokens}
ggml_tensor * indexer_k_pe =
ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens,
ggml_row_size(indexer_k->type, n_embd_indexer_head),
ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0);
cb(indexer_k_pe, "indexer_k_pe", il);
// and {n_embd_indexer_head_nope, 1, n_tokens}
ggml_tensor * indexer_k_nope =
ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens,
ggml_row_size(indexer_k->type, n_embd_indexer_head),
ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1,
ggml_row_size(indexer_k->type, n_embd_indexer_head_nope));
cb(indexer_k_nope, "indexer_k_nope", il);
indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot,
// {n_embd_indexer_head, 1, n_tokens}
indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens);
indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot,
LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(indexer_k_pe, "indexer_k_pe", il);
// {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens}
indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0);
cb(indexer_k, "indexer_k", il);
// perform Hadamard transform on indexer q and k

57
examples/talk-llama/models/dflash.cpp
View File

@ -43,6 +43,8 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);
GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero
if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
}
@ -83,6 +85,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
// reduced draft vocab (optional): d2t maps draft rows to target token ids
int64_t n_vocab_draft = n_vocab;
const struct ggml_tensor * d2t_meta = ml->get_tensor_meta("d2t");
if (d2t_meta) {
n_vocab_draft = d2t_meta->ne[0];
d2t = create_tensor(tn(LLM_TENSOR_D2T), { n_vocab_draft }, 0);
LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft);
}
// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
//
// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
@ -92,7 +104,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
const int64_t dspark_markov_rank = markov_meta->ne[0];
dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0);
dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);
dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0);
dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);
@ -155,6 +167,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
return;
}
// optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
@ -240,6 +255,11 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
const int64_t block_size = std::stoi(it->second);
GGML_ASSERT(block_size > 0);
// bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot
const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor");
const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true";
const int64_t i_draft_beg = sample_from_anchor ? 0 : 1;
const int64_t n_blocks = g.ubatch.n_seqs_unq;
GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks");
// runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size
@ -261,11 +281,26 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
ggml_tensor * cat = nullptr;
ggml_tensor * cat_conf = nullptr;
if (!sample_from_anchor) {
// bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column
cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));
cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));
}
// TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
// token pick, not the Markov conditioning path
for (int64_t i = 0; i < block_drafts; ++i) {
for (int64_t i = i_draft_beg; i < block_drafts; ++i) {
ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab, n_blocks]
ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks]
if (model.d2t) {
// reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)
const int64_t n_draft_vocab = bias->ne[0];
ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f);
bias = ggml_set_rows(ctx0, full,
ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks),
ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1));
bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks);
}
// position i of every block: strided view [n_vocab, n_blocks]
ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]);
@ -495,6 +530,22 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
}
cur = build_lora_mm(output, cur, output_s);
// reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t
if (model.d2t) {
const int64_t n_draft_vocab = cur->ne[0];
const int64_t n_outputs = cur->ne[1];
const int64_t n_vocab = (int64_t) model.vocab.n_tokens();
GGML_ASSERT(model.d2t->type == GGML_TYPE_I64);
GGML_ASSERT(model.d2t->ne[0] == n_draft_vocab);
ggml_tensor * logits = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_outputs), -INFINITY);
cur = ggml_set_rows(ctx0, logits,
ggml_reshape_3d(ctx0, cur, 1, n_draft_vocab, n_outputs),
ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1));
cur = ggml_reshape_2d(ctx0, cur, n_vocab, n_outputs);
}
cb(cur, "result_output", -1);
res->t_logits = cur;

51
examples/talk-llama/models/glm-dsa.cpp
View File

@ -216,10 +216,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
const int64_t n_indexer_head = hparams.indexer_n_head;
const int64_t n_embd_indexer_head = hparams.indexer_head_size;
const int64_t n_embd_indexer_head_rope = hparams.n_rot();
const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope;
const uint32_t n_indexer_top_k = hparams.indexer_top_k;
// the indexer head layout is [rope | nope]
GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head);
const uint32_t kv_lora_rank = hparams.n_lora_kv;
// We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
@ -273,28 +274,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr);
cb(indexer_q, "indexer_q", il);
// split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens}
ggml_tensor * indexer_q_pe =
ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens,
ggml_row_size(indexer_q->type, n_embd_indexer_head),
ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0);
cb(indexer_q_pe, "indexer_q_pe", il);
// and {n_embd_indexer_head_nope, n_indexer_head, n_tokens}
ggml_tensor * indexer_q_nope =
ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens,
ggml_row_size(indexer_q->type, n_embd_indexer_head),
ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head,
ggml_row_size(indexer_q->type, n_embd_indexer_head_nope));
cb(indexer_q_nope, "indexer_q_nope", il);
indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot,
// {n_embd_indexer_head, n_indexer_head, n_tokens}
indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens);
indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot,
LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(indexer_q_pe, "indexer_q_pe", il);
// {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens}
indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0);
cb(indexer_q, "indexer_q", il);
ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur);
@ -303,28 +287,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il);
cb(indexer_k, "indexer_k", il);
// split into {n_embd_indexer_head_rope, 1, n_tokens}
ggml_tensor * indexer_k_pe =
ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens,
ggml_row_size(indexer_k->type, n_embd_indexer_head),
ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0);
cb(indexer_k_pe, "indexer_k_pe", il);
// and {n_embd_indexer_head_nope, 1, n_tokens}
ggml_tensor * indexer_k_nope =
ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens,
ggml_row_size(indexer_k->type, n_embd_indexer_head),
ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1,
ggml_row_size(indexer_k->type, n_embd_indexer_head_nope));
cb(indexer_k_nope, "indexer_k_nope", il);
indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot,
// {n_embd_indexer_head, 1, n_tokens}
indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens);
indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot,
LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(indexer_k_pe, "indexer_k_pe", il);
// {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens}
indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0);
cb(indexer_k, "indexer_k", il);
// perform Hadamard transform on indexer q and k

614
examples/talk-llama/models/kimi-k3.cpp Normal file
View File

@ -0,0 +1,614 @@
#include "models.h"
#include "llama-memory-recurrent.h"
//
// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear.
// Parts that kimi-linear does not have:
// 1. cross-layer residual attention (attn_res_block_size)
// 2. latent MoE (routed experts run at n_expert_latent)
// 3. situ activation (replaces SwiGLU everywhere)
// 4. MLA output gate (sigmoid gate before o_proj)
// 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b)
//
void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false);
ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false);
// the MLA cache holds the compressed latent
// set it here too, as older GGUFs have no value_length key
hparams.n_embd_head_v_full = hparams.n_lora_kv;
// n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear
for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
}
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false);
ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size);
ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta);
ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta);
switch (hparams.n_layer()) {
case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
if (hparams.attn_res_block_size > 0) {
output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 0);
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
if (hparams.attn_res_block_size > 0) {
layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0);
layer.ffn_res_score = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE, "weight", i), {n_embd}, 0);
}
const int64_t head_dim = hparams.n_embd_head_kda;
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = head_dim * n_head;
if (hparams.is_recr(i)) {
// conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1)
auto conv = [&](llm_tensor tid) {
ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);
return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);
};
layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);
layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);
layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);
create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);
layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0);
layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0);
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
// K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded)
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
// K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair
layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0);
layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);
} else {
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t kv_lora_rank = hparams.n_lora_kv;
const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();
const int64_t qk_rope_head_dim = hparams.n_rot();
const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim;
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED);
layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);
if (layer.attn_q_a_norm) {
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0);
} else {
layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0);
}
layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0);
layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i),
{kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)},
TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
if (!layer.wkv_b) {
layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0);
layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0);
}
// K3: sigmoid output gate applied to the attention output before o_proj
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0);
}
if (i < (int) hparams.n_layer_dense_lead) {
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
} else {
const int64_t n_ff_exp = hparams.n_ff_exp;
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
// routed experts live in the latent space
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_latent, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);
if (hparams.n_expert_latent > 0) {
layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0);
layer.ffn_routed_up = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP, "weight", i), {n_embd_latent, n_embd}, 0);
layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED);
}
// shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared
const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
}
}
}
std::unique_ptr<llm_graph_context> llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta)
// linear_beta <= 0 disables the transform on the up branch
static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up,
float beta, float linear_beta) {
ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta);
a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate));
if (linear_beta > 0.0f) {
up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta);
}
return ggml_mul(ctx0, a, up);
}
//
// cross-layer residual attention
//
// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1
// append the new checkpoint, do not re-fold the whole chain
void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) {
ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt;
}
ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w,
int64_t n_tokens, int il) {
if (!resi_stack) {
return cur; // layer 0: nothing banked yet
}
const int n_ckpt = (int) resi_stack->ne[1];
const float eps = hparams.f_norm_rms_eps;
ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens]
// one rms_norm scores all checkpoints at once
// note: the scores use the normalized values, but the sum below uses the raw ones
ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps);
sc_src = ggml_mul(ctx0, sc_src, score_w);
sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens]
sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens);
// the current residual stream is scored apart, so the stack stays append-only
ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps);
sc_cur = ggml_mul(ctx0, sc_cur, score_w);
sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens]
ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0); // [n_ckpt+1, n_tokens]
ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1
cb(probs, "res_probs", il);
// split the sum: hc_pre handles the stack, a broadcast-multiply the current stream
ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0));
ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1],
probs->nb[0] * n_ckpt));
ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src);
out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur));
return out;
}
llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_build_delta_net_base(params), model(model) {
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
cb(inpL, "inp_embd", -1);
// K3 MLA is nope-only, so there is no position input
auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr;
auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr;
auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr();
auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr;
auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr;
ggml_tensor * inp_out_ids = build_inp_out_ids();
const int64_t n_head_kda = hparams.n_head();
const int64_t head_dim = hparams.n_embd_head_kda;
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = n_head_kda * head_dim;
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
GGML_ASSERT(n_seqs != 0);
GGML_ASSERT(ubatch.equal_seqs());
GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
const int64_t kv_lora_rank = hparams.n_lora_kv;
const int64_t n_embd_head_qk_rope = hparams.n_rot();
const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;
const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla);
const uint32_t res_bs = hparams.attn_res_block_size;
const bool use_attn_res = res_bs > 0;
const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
// the residual stream, banked on checkpoint layers and then restarted
// from the attention output alone
ggml_tensor * prefix_sum = inpL;
cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il)
: prefix_sum;
bool banked = false;
if (use_attn_res && (uint32_t) il % res_bs == 0) {
res_push(prefix_sum, n_embd, n_tokens); // banks the RAW layer input, not `cur`
banked = true;
}
cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
ggml_build_forward_expand(gf, cur);
if (hparams.is_recr(il)) {
cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda,
d_inner, n_seq_tokens, n_seqs, il);
} else {
cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv,
n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank,
n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il);
}
prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur);
cb(prefix_sum, "prefix_sum_attn", il);
cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il)
: prefix_sum;
cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
if ((uint32_t) il < hparams.n_layer_dense_lead) {
ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur);
cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);
cur = ggml_mul_mat(ctx0, layer.ffn_down, cur);
cb(cur, "ffn_out", il);
} else {
cur = build_latent_moe(cur, layer, n_embd_latent, il);
}
prefix_sum = ggml_add(ctx0, prefix_sum, cur);
prefix_sum = build_cvec(prefix_sum, il);
cb(prefix_sum, "l_out", il);
inpL = prefix_sum;
}
cur = inpL;
// final mix, then narrow to the output tokens
if (use_attn_res) {
cur = res_mix(cur, model.output_res_score, n_tokens, -1);
}
if (inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = ggml_mul_mat(ctx0, model.output, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
//
// KDA layer
//
// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use
static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
int64_t d_conv, int64_t head_dim, int64_t n_head,
int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) {
const int64_t d_inner = head_dim * n_head;
const int64_t conv_state_size = (d_conv - 1) * d_inner;
const int64_t n_embd_r_total = 3 * conv_state_size;
ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
(d_conv - 1) * ggml_element_size(conv_state_all),
n_embd_r_total * ggml_element_size(conv_state_all),
qkv * conv_state_size * ggml_element_size(conv_state_all));
ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, last_conv_x,
ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
(d_conv - 1) * ggml_element_size(conv_states_all),
n_embd_r_total * ggml_element_size(conv_states_all),
(kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
Xcur = ggml_silu(ctx0, Xcur);
return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
}
ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs,
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
const auto * mctx_cur = inp_rs->mctx;
const auto kv_head = mctx_cur->get_head();
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
cb(Qcur, "kda_q_conv", il);
cb(Kcur, "kda_k_conv", il);
cb(Vcur, "kda_v_conv", il);
// gate_lower_bound is not a clamp - when set, it swaps the decay gate activation:
// unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias)
// set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias))
// ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a
ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a);
g1 = ggml_add(ctx0, g1, layer.ssm_dt_b);
ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1);
if (hparams.kda_gate_lower_bound > -INFINITY) {
g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
g1 = ggml_mul(ctx0, g1, A); // -exp(A_log) * (...)
g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f));
g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound);
} else {
g1 = ggml_softplus(ctx0, g1);
g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
g1 = ggml_mul(ctx0, g1, A);
}
cb(g1, "kda_g1", il);
g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs);
ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs);
beta = ggml_sigmoid(ctx0, beta);
cb(beta, "kda_beta", il);
ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs);
state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs);
const float eps = hparams.f_norm_rms_eps;
Qcur = ggml_l2_norm(ctx0, Qcur, eps);
Kcur = ggml_l2_norm(ctx0, Kcur, eps);
auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il);
ggml_tensor * output = ggml_cont(ctx0, attn_out.first);
cb(output, "kda_scan_out", il);
ggml_tensor * new_state = attn_out.second;
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, new_state,
ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
// K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x)))
ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs);
ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d);
g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs);
ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs);
ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
cb(g2, "kda_g2", il);
cb(normed, "kda_normed", il);
ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2));
gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens);
cur = ggml_mul_mat(ctx0, layer.wo, gated);
cb(cur, "kda_out", il);
return cur;
}
//
// MLA layer (nope-only, with K3's sigmoid output gate)
//
ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer(
ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv,
int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank,
int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) {
ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input
ggml_tensor * Qcur;
if (layer.wq_a) {
Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur);
Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur);
} else {
Qcur = ggml_mul_mat(ctx0, layer.wq, cur);
}
ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
// no RoPE: mla_use_nope is asserted at conversion time
kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
ggml_tensor * out;
if (layer.wk_b && layer.wv_b) {
ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens,
ggml_row_size(Qcur->type, n_embd_head_k_mla),
ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0);
ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens,
ggml_row_size(Qcur->type, n_embd_head_k_mla),
ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head,
ggml_row_size(Qcur->type, n_embd_head_qk_nope));
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0);
ggml_tensor * V = kv_cmpr_3d;
// wo == NULL: the output projection is applied after the gate below
out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il);
} else {
ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens);
ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr);
const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla;
ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,
ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0);
ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens,
ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head),
ggml_row_size(kv->type, n_embd_head_qk_nope)));
ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens);
ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0);
out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il);
}
// K3: attn_output *= sigmoid(g_proj(x)), then o_proj
if (layer.wqkv_gate) {
ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate));
out = ggml_mul(ctx0, out, g);
cb(out, "mla_gated", il);
}
out = ggml_mul_mat(ctx0, layer.wo, out);
cb(out, "mla_out", il);
return out;
}
//
// latent MoE: down-project, run the routed experts in the latent space, norm, up-project;
// shared experts stay at n_embd and read the un-projected input.
//
ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe(
ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) {
ggml_tensor * identity = cur;
ggml_tensor * routed_in = layer.ffn_routed_down
? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur)
: cur;
// the router scores the full-width input while the experts take the latent one,
// so the logits are computed here and passed to build_moe_ffn
ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity);
cb(logits, "ffn_moe_logits", il);
ggml_tensor * moe_out = build_moe_ffn(routed_in,
nullptr, // gate_inp unused: the logits above are passed instead
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
hparams.n_expert,
hparams.n_expert_used,
LLM_FFN_SITU, hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il,
logits);
cb(moe_out, "ffn_moe_out", il);
if (layer.ffn_routed_norm) {
moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il);
}
if (layer.ffn_routed_up) {
moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out);
}
GGML_UNUSED(n_embd_latent);
if (layer.ffn_gate_shexp) {
ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity);
ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp, identity);
ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);
sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh);
cb(sh, "ffn_shexp", il);
moe_out = ggml_add(ctx0, moe_out, sh);
}
cb(moe_out, "ffn_out", il);
return moe_out;
}

520
examples/talk-llama/models/minimax-01.cpp Normal file
View File

@ -0,0 +1,520 @@
#include "models.h"
#include "llama-memory-recurrent.h"
void llama_model_minimax_01::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale);
// we use n_embd_head_la to set recurrent memory n_embd_s
hparams.n_embd_head_la = hparams.n_embd_head_k_full;
// Mark recurrent layers (lightning attention layers).
if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 8;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
}
}
switch (hparams.n_layer()) {
case 80: type = LLM_TYPE_456B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_minimax_01::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
// output
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
// if output is NULL, init from the input tok embed
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
if (!hparams.is_recr(i)) {
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);
} else {
layer.attn_norm_2 = create_tensor(tn(LLM_TENSOR_ATTN_NORM_2, "weight", i), {n_embd_head_k * n_head}, 0);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd_head_k * n_head}, 0);
layer.wg = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);
}
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff, n_expert}, TENSOR_NOT_REQUIRED);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_minimax_01::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
class llm_graph_input_la : public llm_graph_input_i {
public:
llm_graph_input_la(const llama_hparams & hparams) : hparams(hparams) {}
void set_input(const llama_ubatch * ubatch) override {
// this operates on assumption that we have an equal ubatch split
const int64_t n_head = hparams.n_head();
const int32_t n_seqs = ubatch->n_seqs;
const int32_t n_seqs_unq = ubatch->n_seqs_unq;
const int32_t n_tokens = ubatch->n_tokens;
const int32_t n_seq_tokens = ubatch->n_seq_tokens;
std::vector<llama_pos> p0(n_seqs_unq);
std::fill(p0.begin(), p0.end(), std::numeric_limits<llama_pos>::max());
// get lowest token position in a ubatch for each stream
for (int i = 0; i < n_tokens; ++i) {
llama_seq_id seq_id = ubatch->seq_id[i][0];
int32_t seq_idx = ubatch->seq_idx[seq_id];
llama_pos pos = ubatch->pos[i];
if (p0[seq_idx] > pos) {
p0[seq_idx] = pos;
}
}
if (inp_slopes) {
GGML_ASSERT(ggml_backend_buffer_is_host(inp_slopes->buffer));
float * data = (float *) inp_slopes->data;
float start = powf(2, -powf(2, -(log2f(n_head) - 3)));
float ratio = start;
for (int h = 0; h < n_head; ++h) {
data[h] = start * powf(ratio, h);
}
}
if (inp_q_decay) {
GGML_ASSERT(ggml_backend_buffer_is_host(inp_q_decay->buffer));
float * slopes = (float *) inp_slopes->data;
float * data = (float *) inp_q_decay->data;
for (int s = 0; s < n_seqs; ++s) {
for (int i = 0; i < n_seq_tokens; ++i) {
llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0];
int32_t seq_idx = ubatch->seq_idx[seq_id];
llama_pos pos = ubatch->pos[s * n_seq_tokens + i];
int pos_rel = pos - p0[seq_idx];
for (int h = 0; h < n_head; ++h) {
data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (pos_rel + 1);
}
}
}
}
if (inp_k_decay) {
GGML_ASSERT(ggml_backend_buffer_is_host(inp_k_decay->buffer));
float * slopes = (float *) inp_slopes->data;
float * data = (float *) inp_k_decay->data;
for (int s = 0; s < n_seqs; ++s) {
for (int i = 0; i < n_seq_tokens; ++i) {
llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0];
int32_t seq_idx = ubatch->seq_idx[seq_id];
llama_pos pos = ubatch->pos[s * n_seq_tokens + i];
int pos_rel = pos - p0[seq_idx];
for (int h = 0; h < n_head; ++h) {
data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (n_seq_tokens - pos_rel - 1);
}
}
}
}
if (inp_diag_decay) {
GGML_ASSERT(ggml_backend_buffer_is_host(inp_diag_decay->buffer));
float * slopes = (float *) inp_slopes->data;
float * data = (float *) inp_diag_decay->data;
for (int s = 0; s < n_seqs; ++s) {
for (int h = 0; h < n_head; ++h) {
for (int j = 0; j < n_seq_tokens; ++j) {
llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + j][0];
int32_t seq_idx = ubatch->seq_idx[seq_id];
llama_pos pos_j = ubatch->pos[s * n_seq_tokens + j];
int pos_rel_j = pos_j - p0[seq_idx];
for (int i = 0; i < n_seq_tokens; ++i) {
llama_pos pos_i = ubatch->pos[s * n_seq_tokens + i];
int pos_rel_i = pos_i - p0[seq_idx];
int index = pos_rel_j - pos_rel_i;
float s_index = index >= 0 ? -slopes[h] * index : -INFINITY;
data[seq_idx * n_head * n_seq_tokens * n_seq_tokens + h * n_seq_tokens * n_seq_tokens + j * n_seq_tokens + i] = s_index;
}
}
}
}
}
}
bool can_reuse(const llm_graph_params & params) override {
bool res = true;
if (params.ubatch.n_seq_tokens > 1) {
res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens);
res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens);
res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens);
}
return res;
}
const llama_hparams & hparams;
ggml_tensor * inp_slopes = nullptr; // F32 [n_head]
ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch]
ggml_tensor * inp_k_decay = nullptr; // F32 [1, n_head, n_batch]
ggml_tensor * inp_diag_decay = nullptr; // F32 [n_batch, n_batch, n_head]
};
llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
// GGML_ASSERT(n_embd_head == n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
GGML_ASSERT(n_seqs != 0);
GGML_ASSERT(ubatch.equal_seqs());
GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
auto * inp_hybrid = build_inp_mem_hybrid();
auto * inp_rs = inp_hybrid->get_recr();
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
llm_graph_input_la * la = nullptr;
auto inp = std::make_unique<llm_graph_input_la>(hparams);
inp->inp_slopes = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_head);
ggml_set_input(inp->inp_slopes);
cb(inp->inp_slopes, "slopes", -1);
if (n_seq_tokens != 1) {
inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs);
ggml_set_input(inp->inp_q_decay);
cb(inp->inp_q_decay, "q_decay_exp", -1);
inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs);
ggml_set_input(inp->inp_k_decay);
cb(inp->inp_k_decay, "k_decay_exp", -1);
inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs);
ggml_set_input(inp->inp_diag_decay);
cb(inp->inp_diag_decay, "diag_decay_exp", -1);
}
la = (llm_graph_input_la *) res->add_input(std::move(inp));
ggml_tensor * slopes = la->inp_slopes;
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
ggml_tensor * inpSA = inpL;
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
ggml_tensor * residual = cur;
// self_attention
if (!hparams.is_recr(il)) {
// softmax attention layer
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
cur = build_attn(inp_hybrid->get_attn(),
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
} else {
// lightning attention layer
const auto * mctx_cur = inp_rs->mctx;
const auto kv_head = mctx_cur->get_head();
// TODO unneeded - any way to make conv states optional in recurrent memory?
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
ggml_build_forward_expand(gf, conv_state_all);
float slope_scale = 1.0 - 1.0 * il / (n_layer - 1) + 1e-5;
ggml_tensor * slope_rate = ggml_scale(ctx0, slopes, slope_scale);
cb(slope_rate, "slope_rate", il);
cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], n_seq_tokens, 1, n_seqs);
ggml_tensor * QKVcur = build_lora_mm(model.layers[il].wqkv, cur);
cb(QKVcur, "QKVcur", il);
QKVcur = ggml_silu(ctx0, QKVcur);
cb(QKVcur, "QKVcur_silu", il);
QKVcur = ggml_reshape_4d(ctx0, QKVcur, n_embd_head * 3, n_head, n_seq_tokens, n_seqs);
ggml_tensor * Qcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 0*ggml_element_size(QKVcur)*n_embd_head);
ggml_tensor * Kcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 1*ggml_element_size(QKVcur)*n_embd_head);
ggml_tensor * Vcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 2*ggml_element_size(QKVcur)*n_embd_head);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
// get previous KV
ggml_tensor * la_states_all = mctx_cur->get_s_l(il);
ggml_tensor * state = build_rs(inp_rs, la_states_all, hparams.n_embd_s(), n_seqs);
ggml_tensor * kv_old = ggml_reshape_4d(ctx0, state, n_embd_head, n_embd_head, n_head, n_seqs);
cb(kv_old, "kv_old", il);
ggml_tensor * qkv = nullptr;
ggml_tensor * kv_new = nullptr;
if (n_seq_tokens == 1) {
// lightning attention - optimized single token case for TG
ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0);
cb(slopes_neg, "slopes_neg", il);
ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg);
cb(ratio, "ratio", il);
ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head);
cb(ratio_3d, "ratio3d", il);
ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3));
cb(v_trans, "v_trans", il);
ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3));
cb(k_trans, "k_trans", il);
ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans);
cb(kv_cur, "kv_cur", il);
ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d);
cb(kv_old_s, "kv_old_s", il);
kv_new = ggml_add(ctx0, kv_old_s, kv_cur);
cb(kv_new, "kv_new", il);
ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
cb(q_trans, "q_trans", il);
qkv = ggml_mul_mat(ctx0, kv_new, q_trans);
cb(qkv, "qkv", il);
} else if(n_seq_tokens > 1) {
// lightning attention - general multi token case for PP
ggml_tensor * q_decay_exp = la->inp_q_decay;
ggml_tensor * k_decay_exp = la->inp_k_decay;
ggml_tensor * diag_decay_exp = la->inp_diag_decay;
ggml_tensor * q_decay = ggml_exp(ctx0, ggml_scale(ctx0, q_decay_exp, slope_scale));
cb(q_decay, "q_decay", il);
ggml_tensor * k_decay = ggml_exp(ctx0, ggml_scale(ctx0, k_decay_exp, slope_scale));
cb(k_decay, "k_decay", il);
ggml_tensor * diag_decay = ggml_exp(ctx0, ggml_scale(ctx0, diag_decay_exp, slope_scale));
cb(diag_decay, "diag_decay", il);
ggml_tensor * q_s = ggml_mul(ctx0, Qcur, q_decay);
cb(q_s, "q_s", il);
ggml_tensor * q_s_trans = ggml_permute(ctx0, q_s, 0, 2, 1, 3);
cb(q_s_trans, "q_s_trans", il);
ggml_tensor * qkv_none_diag = ggml_mul_mat(ctx0, kv_old, q_s_trans);
cb(qkv_none_diag, "qkv_none_diag", il);
ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
cb(q_trans, "q_trans", il);
ggml_tensor * k_trans = ggml_permute(ctx0, Kcur, 0, 2, 1, 3);
cb(k_trans, "k_trans", il);
ggml_tensor * qk = ggml_mul_mat(ctx0, k_trans, q_trans);
cb(qk, "qk", il);
qk = ggml_mul(ctx0, qk, diag_decay);
cb(qk, "qk_s", il);
ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3));
cb(v_trans, "v_trans", il);
ggml_tensor * qkv_diag = ggml_mul_mat(ctx0, v_trans, qk);
cb(qkv_diag, "qkv_diag", il);
qkv = ggml_add(ctx0, qkv_none_diag, qkv_diag);
cb(qkv, "qkv", il);
ggml_build_forward_expand(gf, qkv);
ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0*n_seq_tokens);
cb(slopes_neg, "slopes_neg", il);
ggml_tensor * block_decay = ggml_exp(ctx0, slopes_neg);
cb(block_decay, "block_decay", il);
ggml_tensor * block_decay_3d = ggml_reshape_3d(ctx0, block_decay, 1, 1, n_head);
cb(block_decay_3d, "block_decay_3d", il);
ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, block_decay_3d);
cb(kv_old_s, "kv_old_s", il);
ggml_tensor * k_after_decay = ggml_mul(ctx0, Kcur, k_decay);
cb(k_after_decay, "k_after_decay", il);
ggml_tensor * k_after_decay_trans = ggml_cont(ctx0, ggml_permute(ctx0, k_after_decay, 1, 2, 0, 3));
cb(k_after_decay_trans, "k_after_decay_trans", il);
ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_after_decay_trans, v_trans);
cb(kv_cur, "kv_cur", il);
kv_new = ggml_add(ctx0, kv_old_s, kv_cur);
cb(kv_new, "kv_new", il);
}
// store new KV
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, kv_new,
ggml_view_1d(ctx0, la_states_all, hparams.n_embd_s() * n_seqs,
kv_head * hparams.n_embd_s() * ggml_element_size(la_states_all))));
qkv = ggml_cont(ctx0, ggml_permute(ctx0, qkv, 0, 2, 1, 3));
cb(qkv, "qkv_permuted", il);
qkv = ggml_reshape_4d(ctx0, qkv, qkv->ne[0]*qkv->ne[1], qkv->ne[2], 1, qkv->ne[3]);
// norm
ggml_tensor * qkv_norm = build_norm(qkv,
model.layers[il].attn_norm_2, NULL,
LLM_NORM_RMS, il);
cb(qkv_norm, "qkv_norm", il);
ggml_tensor * g = build_lora_mm(model.layers[il].wg, cur);
cb(g, "g", il);
g = ggml_sigmoid(ctx0, g);
cb(g, "g_sigm", il);
cur = ggml_mul(ctx0, g, qkv_norm);
cur = build_lora_mm(model.layers[il].wo, cur);
cb(cur, "attn_out", il);
cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens*n_seqs);
cb(cur, "attn_out", il);
}
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
residual = ggml_get_rows(ctx0, residual, inp_out_ids);
}
residual = ggml_scale(ctx0, residual, hparams.f_residual_scale);
cb(residual, "residual_scaled_attn", il);
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual);
cb(ffn_inp, "ffn_inp", il);
// MoE branch
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm, NULL,
LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
residual = cur;
cur = build_moe_ffn(cur,
model.layers[il].ffn_gate_inp,
model.layers[il].ffn_up_exps,
model.layers[il].ffn_gate_exps,
model.layers[il].ffn_down_exps,
model.layers[il].ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_SILU, true,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,
il);
cb(cur, "ffn_moe_out", il);
residual = ggml_scale(ctx0, residual, hparams.f_residual_scale);
cb(residual, "residual_scaled_ffn", il);
cur = ggml_add(ctx0, cur, residual);
cb(cur, "ffn_out", il);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
// input for next layer
inpL = cur;
}
cur = inpL;
cur = build_norm(cur,
model.output_norm, NULL,
LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
// lm_head
cur = build_lora_mm(model.output, cur, model.output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}

2
examples/talk-llama/models/minimax-m3.cpp
View File

@ -25,6 +25,8 @@ void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks };
GGML_ASSERT(hparams.indexer_block_size > 0); // avoid div by zero
switch (hparams.n_layer()) {
case 60: type = LLM_TYPE_428B_A23B; break;
default: type = LLM_TYPE_UNKNOWN;

68
examples/talk-llama/models/models.h
View File

@ -1784,6 +1784,25 @@ struct llama_model_bailingmoe2 : public llama_model_base {
};
struct llama_model_bailingmoe3 : public llama_model_base {
llama_model_bailingmoe3(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_build_delta_net_base {
graph(const llama_model & model, const llm_graph_params & params);
const llama_model & model;
};
struct graph_mtp : public llm_graph_context {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_seed_oss : public llama_model_base {
llama_model_seed_oss(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@ -2043,6 +2062,19 @@ struct llama_model_apertus : public llama_model_base {
};
struct llama_model_minimax_01 : public llama_model_base {
llama_model_minimax_01(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_minimax_m2 : public llama_model_base {
llama_model_minimax_m2(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
@ -2272,6 +2304,42 @@ struct llama_model_mimo2 : public llama_model_base {
};
struct llama_model_kimi_k3 : public llama_model_base {
llama_model_kimi_k3(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_build_delta_net_base {
graph(const llama_model & model, const llm_graph_params & params);
const llama_model & model;
// Cross-layer residual attention (K3's `_apply_attn_res`).
ggml_tensor * resi_stack = nullptr;
void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens);
ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w,
int64_t n_tokens, int il);
ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_rs * inp_rs,
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_attn_k * inp_attn_k,
llm_graph_input_attn_kv * inp_attn_kv,
int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla,
int64_t kv_lora_rank, int64_t n_embd_head_qk_rope,
int64_t n_embd_head_qk_nope, float kq_scale, int il);
ggml_tensor * build_latent_moe(ggml_tensor * cur, const llama_layer & layer,
int64_t n_embd_latent, int il);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_kimi_linear : public llama_model_base {
llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;