model: add GLM 5.2 Indexer support (#25407)
* Start building graph - reuse deepseek32 * Enable kv cache and rotation for glm_dsa architecture Just follow Deepseek 3.2 for now. * Reuse prev_top_k for "shared" indexer layers * GLM 5.2 uses LLAMA_ROPE_TYPE_NORM for the indexer. This is transformers' `apply_rotary_pos_emb_interleave` * Default indexer types to GLM pattern Previous converted GGUFs like https://huggingface.co/unsloth/GLM-5.2-GGUF write indexer weights to _all_ layers, even if they are only required for "full" types. This PR relies on a new key "%s.attention.indexer.types"; if absent, it will use the default GLM 5.2 schedule as defined in https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L26. Note that conversion is not saving this key yet. * Save indexer types to gguf, restore on load * Use ggml_lightning_indexer when cparams.fused_lid Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com> * GLM 5 and 5.1 use full indexers Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com> * Fix indentation * Ensure array is zero-filled * Prefer explicit std::fill * Assert prev_top_k exists for shared indexer --------- Co-authored-by: fairydreaming <166155368+fairydreaming@users.noreply.github.com>
This commit is contained in:
Pedro Cuenca
co-authored by
fairydreaming
fairydreaming
parent
95a923a64c
commit
88bfee1429
@@ -237,6 +237,9 @@ class GlmMoeDsaModel(DeepseekV2Model):
|
||||
self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
|
||||
self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
|
||||
self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
|
||||
if (indexer_types := self.hparams.get("indexer_types")) is not None:
|
||||
indexer_types = [t == "full" for t in indexer_types]
|
||||
self.gguf_writer.add_indexer_types(indexer_types)
|
||||
|
||||
|
||||
@ModelBase.register("SolarOpenForCausalLM")
|
||||
|
||||
@@ -200,6 +200,7 @@ class Keys:
|
||||
HEAD_COUNT = "{arch}.attention.indexer.head_count"
|
||||
KEY_LENGTH = "{arch}.attention.indexer.key_length"
|
||||
TOP_K = "{arch}.attention.indexer.top_k"
|
||||
TYPES = "{arch}.attention.indexer.types"
|
||||
|
||||
class HyperConnection:
|
||||
COUNT = "{arch}.hyper_connection.count"
|
||||
|
||||
@@ -793,6 +793,10 @@ class GGUFWriter:
|
||||
def add_indexer_top_k(self, top_k: int) -> None:
|
||||
self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k)
|
||||
|
||||
def add_indexer_types(self, value: Sequence[bool]) -> None:
|
||||
key = Keys.Attention.Indexer.TYPES.format(arch=self.arch)
|
||||
self.add_array(key, value)
|
||||
|
||||
def add_max_alibi_bias(self, bias: float) -> None:
|
||||
self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
|
||||
|
||||
|
||||
@@ -253,6 +253,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
|
||||
{ LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" },
|
||||
{ LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" },
|
||||
{ LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" },
|
||||
{ LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" },
|
||||
|
||||
@@ -258,6 +258,7 @@ enum llm_kv {
|
||||
LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,
|
||||
LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,
|
||||
LLM_KV_ATTENTION_INDEXER_TOP_K,
|
||||
LLM_KV_ATTENTION_INDEXER_TYPES,
|
||||
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
|
||||
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
|
||||
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
|
||||
|
||||
@@ -248,6 +248,14 @@ bool llama_hparams::is_mla() const {
|
||||
return n_embd_head_k_mla_impl != 0 && n_embd_head_v_mla_impl != 0;
|
||||
}
|
||||
|
||||
bool llama_hparams::is_indexer_full(uint32_t il) const {
|
||||
if (il < n_layer()) {
|
||||
return is_indexer_full_impl[il];
|
||||
}
|
||||
|
||||
GGML_ABORT("%s: il (%u) out of bounds (n_layer: %u)\n", __func__, il, n_layer());
|
||||
}
|
||||
|
||||
uint32_t llama_hparams::n_embd_head_k_mla() const {
|
||||
return is_mla() ? n_embd_head_k_mla_impl : n_embd_head_k();
|
||||
}
|
||||
|
||||
@@ -227,6 +227,10 @@ struct llama_hparams {
|
||||
uint32_t indexer_head_size = 0;
|
||||
uint32_t indexer_top_k = 0;
|
||||
|
||||
// Indexer is "full" (1) or "shared" (0)
|
||||
// Shared indexers reuse top-k from previous full layer
|
||||
std::array<uint32_t, LLAMA_MAX_LAYERS> is_indexer_full_impl;
|
||||
|
||||
// DeepSeek-V4
|
||||
uint32_t dsv4_o_group_count = 0;
|
||||
uint32_t dsv4_o_lora_rank = 0;
|
||||
@@ -302,6 +306,8 @@ struct llama_hparams {
|
||||
|
||||
bool is_swa(uint32_t il) const;
|
||||
|
||||
bool is_indexer_full(uint32_t il) const;
|
||||
|
||||
void set_recr_pattern(uint32_t n_pattern, bool dense_first = false);
|
||||
|
||||
// whether or not the given layer is recurrent (for hybrid models)
|
||||
|
||||
@@ -323,7 +323,7 @@ llama_kv_cache::llama_kv_cache(
|
||||
hparams.n_embd_head_k() % 64 == 0;
|
||||
|
||||
// always create Hadamard rotation tensors for DeepSeek lightning indexers
|
||||
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4) &&
|
||||
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || model.arch == LLM_ARCH_GLM_DSA) &&
|
||||
hparams.n_embd_head_k_full == hparams.indexer_head_size) {
|
||||
attn_rot_k = true;
|
||||
}
|
||||
|
||||
@@ -281,6 +281,7 @@ void llama_model_saver::add_kv_from_model() {
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
|
||||
add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true);
|
||||
add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true);
|
||||
|
||||
const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train;
|
||||
|
||||
@@ -1129,6 +1129,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
|
||||
std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0);
|
||||
std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0);
|
||||
std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0);
|
||||
std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0);
|
||||
|
||||
std::fill(hparams.xielu_alpha_n.begin(), hparams.xielu_alpha_n.end(), 0.0f);
|
||||
std::fill(hparams.xielu_alpha_p.begin(), hparams.xielu_alpha_p.end(), 0.0f);
|
||||
@@ -2065,6 +2066,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
|
||||
res = nullptr;
|
||||
} break;
|
||||
case LLM_ARCH_DEEPSEEK32:
|
||||
case LLM_ARCH_GLM_DSA:
|
||||
{
|
||||
res = new llama_kv_cache_dsa(
|
||||
*this,
|
||||
|
||||
+395
-2
@@ -1,5 +1,31 @@
|
||||
#include "models.h"
|
||||
|
||||
#include "llama-kv-cache-dsa.h"
|
||||
|
||||
// https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L26
|
||||
const std::array<uint32_t, LLAMA_MAX_LAYERS> GLM_5_2_DEFAULT_INDEXER_TYPES = {
|
||||
1, 1,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
1, 0, 0, 0,
|
||||
};
|
||||
|
||||
void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) {
|
||||
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
|
||||
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
||||
@@ -34,10 +60,19 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) {
|
||||
|
||||
// NextN/MTP parameters
|
||||
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
|
||||
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_impl");
|
||||
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
|
||||
|
||||
// BC for GLM 5, 5.1 (full indexers) without indexer_types metadata
|
||||
const bool is_pre_5_2 = hparams.n_ctx_train < 1048576;
|
||||
if (is_pre_5_2) {
|
||||
std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1);
|
||||
} else {
|
||||
hparams.is_indexer_full_impl = GLM_5_2_DEFAULT_INDEXER_TYPES;
|
||||
}
|
||||
ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);
|
||||
|
||||
switch (hparams.n_layer()) {
|
||||
case 79: type = LLM_TYPE_744B_A40B; break;
|
||||
case 78: type = LLM_TYPE_744B_A40B; break;
|
||||
default: type = LLM_TYPE_UNKNOWN;
|
||||
}
|
||||
}
|
||||
@@ -150,3 +185,361 @@ std::unique_ptr<llm_graph_context> llama_model_glm_dsa::build_arch_graph(const l
|
||||
return std::make_unique<graph>(*this, params);
|
||||
}
|
||||
|
||||
llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_params & params) :
|
||||
llm_graph_context(params) {
|
||||
const bool is_mla = hparams.is_mla();
|
||||
GGML_ASSERT(is_mla);
|
||||
|
||||
// note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA
|
||||
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();
|
||||
GGML_UNUSED(n_embd_head_v);
|
||||
|
||||
const int64_t n_embd_head_qk_rope = hparams.n_rot();
|
||||
const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;
|
||||
|
||||
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;
|
||||
|
||||
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.
|
||||
// See https://github.com/ggml-org/llama.cpp/discussions/7416 for detailed explanation.
|
||||
// And also: https://github.com/ggml-org/llama.cpp/pull/17945 [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]
|
||||
|
||||
// first cancel the adjustment from llama_hparams::yarn_attn_factor_adjust to get the original attn_factor
|
||||
GGML_ASSERT(ext_factor >= 0.0f);
|
||||
const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));
|
||||
|
||||
// use the original attn_factor to pre-scale the kq_scale
|
||||
const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
|
||||
const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));
|
||||
|
||||
ggml_tensor * cur;
|
||||
ggml_tensor * inpL;
|
||||
|
||||
// {n_embd, n_tokens}
|
||||
inpL = build_inp_embd(model.tok_embd);
|
||||
|
||||
// inp_pos - contains the positions
|
||||
ggml_tensor * inp_pos = build_inp_pos();
|
||||
|
||||
llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa();
|
||||
|
||||
ggml_tensor * inp_out_ids = build_inp_out_ids();
|
||||
|
||||
// Difference vs Deepseek 3.2: shared indexer layers reuse the top_k from the previous full indexer layers
|
||||
// See https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L30
|
||||
ggml_tensor * prev_top_k = nullptr;
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
ggml_tensor * inpSA = inpL;
|
||||
|
||||
// norm
|
||||
cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "attn_norm", il);
|
||||
|
||||
// self_attention
|
||||
{
|
||||
ggml_tensor * qr = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur);
|
||||
cb(qr, "qr", il);
|
||||
|
||||
qr = build_norm(qr, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(qr, "qr", il);
|
||||
|
||||
ggml_tensor * top_k = nullptr;
|
||||
|
||||
// lightning indexer
|
||||
if (hparams.is_indexer_full(il)) {
|
||||
// "full" layer
|
||||
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,
|
||||
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);
|
||||
cb(indexer_k, "indexer_k", il);
|
||||
|
||||
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,
|
||||
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
|
||||
indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q);
|
||||
cb(indexer_q, "indexer_q", il);
|
||||
indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k);
|
||||
cb(indexer_k, "indexer_k", il);
|
||||
|
||||
// store indexer keys to KV cache
|
||||
const auto * mctx_lid = inp_attn_dsa->mctx->get_lid();
|
||||
const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid();
|
||||
ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il));
|
||||
|
||||
// prepare indexer weights
|
||||
ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur);
|
||||
cb(indexer_weights, "indexer_weights", il);
|
||||
|
||||
// get cached indexer keys
|
||||
indexer_k = mctx_lid->get_k(ctx0, il);
|
||||
|
||||
// split the batch into streams if needed
|
||||
const auto n_stream = indexer_k->ne[3];
|
||||
indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);
|
||||
indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
|
||||
|
||||
// pre-scale weights to avoid scaling operations on huge indexer_score tensor
|
||||
indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));
|
||||
cb(indexer_weights, "indexer_weights", il);
|
||||
|
||||
ggml_tensor * indexer_score = nullptr;
|
||||
if (cparams.fused_lid) {
|
||||
indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
|
||||
} else {
|
||||
// calculate indexer kq
|
||||
indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
|
||||
cb(indexer_q, "indexer_q", il);
|
||||
indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
|
||||
cb(indexer_k, "indexer_k", il);
|
||||
|
||||
ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
|
||||
cb(indexer_kq, "indexer_kq", il);
|
||||
|
||||
// ReLU requires contiguous tensors
|
||||
indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
|
||||
cb(indexer_kq, "indexer_kq", il);
|
||||
|
||||
// apply ReLU
|
||||
indexer_score = ggml_relu(ctx0, indexer_kq);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// multiply scores by indexer weights
|
||||
indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// sum by q n_indexer_head dimension
|
||||
indexer_score = ggml_sum_rows(ctx0, indexer_score);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// permute result to match KQ mask
|
||||
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
|
||||
// mask indexer scores
|
||||
ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
|
||||
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
|
||||
cb(indexer_score, "indexer_score", il);
|
||||
}
|
||||
|
||||
// get indices of top k indexer scores
|
||||
uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;
|
||||
top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));
|
||||
prev_top_k = top_k;
|
||||
cb(top_k, "top_k", il);
|
||||
} else {
|
||||
// "shared" indexer layer - reuse top-k from a previous full layer
|
||||
GGML_ASSERT(prev_top_k != nullptr && "shared indexer layer must follow a previous full indexer layer");
|
||||
top_k = prev_top_k;
|
||||
cb(top_k, "top_k", il);
|
||||
}
|
||||
|
||||
ggml_tensor * q = ggml_mul_mat(ctx0, model.layers[il].wq_b, qr);
|
||||
cb(q, "q", il);
|
||||
|
||||
// split into {n_embd_head_qk_nope, n_head, n_tokens}
|
||||
ggml_tensor * q_nope =
|
||||
ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
|
||||
ggml_row_size(q->type, n_embd_head_k) * n_head, 0);
|
||||
cb(q_nope, "q_nope", il);
|
||||
|
||||
// and {n_embd_head_qk_rope, n_head, n_tokens}
|
||||
ggml_tensor * q_pe = ggml_view_3d(
|
||||
ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
|
||||
ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));
|
||||
cb(q_pe, "q_pe", il);
|
||||
|
||||
ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);
|
||||
cb(kv_cmpr_pe, "kv_cmpr_pe", il);
|
||||
|
||||
// split into {kv_lora_rank, n_tokens}
|
||||
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);
|
||||
cb(kv_cmpr, "kv_cmpr", il);
|
||||
|
||||
// and {n_embd_head_qk_rope, 1, n_tokens}
|
||||
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));
|
||||
cb(k_pe, "k_pe", il);
|
||||
|
||||
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);
|
||||
cb(q_pe, "q_pe", il);
|
||||
|
||||
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);
|
||||
cb(k_pe, "k_pe", il);
|
||||
|
||||
kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
|
||||
cb(kv_cmpr, "kv_cmpr", il);
|
||||
|
||||
// MLA attention
|
||||
{
|
||||
// {n_embd_head_qk_nope, n_tokens, n_head}
|
||||
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
|
||||
cb(q_nope, "q_nope_perm", il);
|
||||
|
||||
// {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}
|
||||
ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, model.layers[il].wk_b, q_nope);
|
||||
cb(q_nope_absorbed, "q_nope_absorbed", il);
|
||||
|
||||
// {kv_lora_rank, n_head, n_tokens}
|
||||
q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
|
||||
cb(q_nope_absorbed, "q_nope_absorbed_perm", il);
|
||||
|
||||
// {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}
|
||||
// note: rope must go first for in-place context shifting in build_rope_shift()
|
||||
ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
|
||||
cb(Qcur, "Qcur", il);
|
||||
|
||||
kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
|
||||
cb(kv_cmpr, "kv_cmpr_reshape", il);
|
||||
|
||||
// {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}
|
||||
ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
|
||||
cb(Kcur, "Kcur", il);
|
||||
|
||||
// {kv_lora_rank, 1, n_tokens}
|
||||
ggml_tensor * Vcur = kv_cmpr;
|
||||
cb(Vcur, "Vcur", il);
|
||||
|
||||
// note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
|
||||
cur = build_attn(inp_attn_dsa,
|
||||
model.layers[il].wo, NULL, model.layers[il].wo_s,
|
||||
Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, 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);
|
||||
}
|
||||
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
|
||||
cb(ffn_inp, "ffn_inp", il);
|
||||
|
||||
cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
if ((uint32_t) il < hparams.n_layer_dense_lead) {
|
||||
cur = build_ffn(cur,
|
||||
model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,
|
||||
model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
|
||||
model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
|
||||
NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(cur, "ffn_out", il);
|
||||
} else {
|
||||
// MoE branch
|
||||
ggml_tensor * moe_out = 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, hparams.expert_weights_norm,
|
||||
hparams.expert_weights_scale,
|
||||
(llama_expert_gating_func_type) hparams.expert_gating_func,
|
||||
il,
|
||||
nullptr,
|
||||
model.layers[il].ffn_gate_up_exps,
|
||||
model.layers[il].ffn_up_exps_s,
|
||||
model.layers[il].ffn_gate_exps_s,
|
||||
model.layers[il].ffn_down_exps_s);
|
||||
cb(moe_out, "ffn_moe_out", il);
|
||||
|
||||
// FFN shared expert
|
||||
{
|
||||
ggml_tensor * ffn_shexp =
|
||||
build_ffn(cur,
|
||||
model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
|
||||
model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
|
||||
model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
|
||||
NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(ffn_shexp, "ffn_shexp", il);
|
||||
|
||||
cur = ggml_add(ctx0, moe_out, ffn_shexp);
|
||||
cb(cur, "ffn_out", il);
|
||||
}
|
||||
}
|
||||
cur = ggml_add(ctx0, cur, ffn_inp);
|
||||
|
||||
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 = ggml_mul_mat(ctx0, model.output, cur);
|
||||
|
||||
cb(cur, "result_output", -1);
|
||||
res->t_logits = cur;
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
+3
-1
@@ -1217,7 +1217,9 @@ struct llama_model_glm_dsa : public llama_model_base {
|
||||
void load_arch_hparams(llama_model_loader & ml) override;
|
||||
void load_arch_tensors(llama_model_loader & ml) override;
|
||||
|
||||
using graph = llama_model_deepseek2::graph;
|
||||
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;
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user