From 25b654c74b96215d8f2053d0f30108da893d4a7d Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 14 Aug 2026 13:21:21 +0200 Subject: [PATCH] init impl --- src/llama-arch.cpp | 5 + src/llama-arch.h | 4 + src/llama-graph.cpp | 73 ++++-- src/llama-graph.h | 25 +- src/llama-hparams.h | 5 + src/llama-kv-cache-dsa.cpp | 107 +++++++- src/llama-kv-cache-dsa.h | 15 +- src/llama-kv-cache.cpp | 3 +- src/llama-model-saver.cpp | 1 + src/llama-model.cpp | 12 +- src/models/dots3note.cpp | 510 +++++++++++++++++++++++++++++++++++++ src/models/models.h | 12 + tests/test-llama-archs.cpp | 16 +- 13 files changed, 740 insertions(+), 48 deletions(-) create mode 100644 src/models/dots3note.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 8ed9391d7c..89c7f75ca1 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -108,6 +108,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_BAILINGMOE, "bailingmoe" }, { LLM_ARCH_BAILINGMOE2, "bailingmoe2" }, { LLM_ARCH_DOTS1, "dots1" }, + { LLM_ARCH_DOTS3NOTE, "dots3note" }, { LLM_ARCH_ARCEE, "arcee" }, { LLM_ARCH_AFMOE, "afmoe" }, { LLM_ARCH_LAGUNA, "laguna" }, @@ -263,6 +264,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_VALUE_LENGTH_MLA, "%s.attention.value_length_mla" }, { LLM_KV_ATTENTION_KEY_LENGTH_SWA, "%s.attention.key_length_swa" }, { LLM_KV_ATTENTION_VALUE_LENGTH_SWA, "%s.attention.value_length_swa" }, + { LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA, "%s.attention.key_length_mla_swa" }, + { LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA, "%s.attention.value_length_mla_swa" }, + { LLM_KV_ATTENTION_KV_LORA_RANK_SWA, "%s.attention.kv_lora_rank_swa" }, { 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" }, @@ -1023,6 +1027,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DOTS3NOTE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_BITNET: case LLM_ARCH_T5: diff --git a/src/llama-arch.h b/src/llama-arch.h index 18d9de186f..34def8880b 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -113,6 +113,7 @@ enum llm_arch { LLM_ARCH_BAILINGMOE, LLM_ARCH_BAILINGMOE2, LLM_ARCH_DOTS1, + LLM_ARCH_DOTS3NOTE, LLM_ARCH_ARCEE, LLM_ARCH_AFMOE, LLM_ARCH_LAGUNA, @@ -268,6 +269,9 @@ enum llm_kv { LLM_KV_ATTENTION_VALUE_LENGTH_MLA, LLM_KV_ATTENTION_KEY_LENGTH_SWA, LLM_KV_ATTENTION_VALUE_LENGTH_SWA, + LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA, + LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA, + LLM_KV_ATTENTION_KV_LORA_RANK_SWA, LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 55d8580246..594b1b60b2 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -559,6 +559,12 @@ void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) { mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn); + if (mctx->get_mla_swa()) { + mctx->get_mla_swa()->set_input_k_idxs(self_k_idxs_mla_swa, ubatch); + + mctx->get_mla_swa()->set_input_kq_mask(self_kq_mask_mla_swa, ubatch, cparams.causal_attn); + } + mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch); mctx->get_lid()->set_input_kq_mask(self_kq_mask_lid, ubatch, cparams.causal_attn); @@ -579,6 +585,11 @@ bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) { res &= can_reuse_kq_mask(self_kq_mask_mla, mctx->get_mla(), params.ubatch, params.cparams); res &= can_reuse_kq_mask(self_kq_mask_lid, mctx->get_lid(), params.ubatch, params.cparams); + if (mctx->get_mla_swa()) { + res &= self_k_idxs_mla_swa && self_k_idxs_mla_swa->ne[0] == params.ubatch.n_tokens; + res &= self_kq_mask_mla_swa && can_reuse_kq_mask(self_kq_mask_mla_swa, mctx->get_mla_swa(), params.ubatch, params.cparams); + } + return res; } @@ -2925,48 +2936,57 @@ ggml_tensor * llm_graph_context::build_attn( ggml_build_forward_expand(gf, v_cur); ggml_build_forward_expand(gf, k_cur); - const auto * mctx_cur = inp->mctx->get_mla(); + const bool is_swa = hparams.is_swa(il); + + // SWA layers use the dedicated window cache and have no indexer, so no top-k selection + GGML_ASSERT(!is_swa || top_k == nullptr); + + const auto * mctx_cur = is_swa ? inp->mctx->get_mla_swa() : inp->mctx->get_mla(); // store to KV cache { - const auto & k_idxs = inp->get_k_idxs_mla(); + const auto & k_idxs = is_swa ? inp->get_k_idxs_mla_swa() : inp->get_k_idxs_mla(); ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); } - const auto & kq_mask = inp->get_kq_mask_mla(); + const auto & kq_mask = is_swa ? inp->get_kq_mask_mla_swa() : inp->get_kq_mask_mla(); - // prepare new kq mask - starts filled with -INFINITY - ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + ggml_tensor * kq_mask_used = kq_mask; - // reshape KQ mask into tensor with rows of size 1: - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] - kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + if (top_k) { + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); - // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] - ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); - // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] - // this will be our source of zero values for unmasking top k mask elements - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); - zeros = ggml_fill(ctx0, zeros, 0.0f); + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - // modify KQ mask by unmasking elements that are in top_k indices - // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) - ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); - // reshape to restore the original shape of KQ mask: - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] - kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); - // combine with the original kq mask - kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + // combine with the original kq mask + kq_mask_used = ggml_add(ctx0, kq_mask_top_k, kq_mask); + } ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_used, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { @@ -3207,6 +3227,13 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const { inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla; } + if (mctx_cur->get_mla_swa()) { + inp->self_k_idxs_mla_swa = mctx_cur->get_mla_swa()->build_input_k_idxs(ctx0, ubatch); + + inp->self_kq_mask_mla_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla_swa(), ubatch, cparams); + inp->self_kq_mask_mla_swa_cnv = inp->self_kq_mask_mla_swa; + } + { inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch); diff --git a/src/llama-graph.h b/src/llama-graph.h index 75bc0fe80d..fbd2ceb271 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -404,19 +404,24 @@ public: bool can_reuse(const llm_graph_params & params) override; - ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; } - ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; } + ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; } + ggml_tensor * get_k_idxs_mla_swa() const { return self_k_idxs_mla_swa; } + ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; } - ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; } - ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; } + ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; } + ggml_tensor * get_kq_mask_mla_swa() const { return self_kq_mask_mla_swa_cnv; } + ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; } - ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch] - ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch] + ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch] + ggml_tensor * self_k_idxs_mla_swa = nullptr; // I64 [n_batch] + ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch] - ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream] - ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_mla_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_mla_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] ggml_tensor * self_k_rot_lid = nullptr; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 57de808242..c5eb2167c5 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -100,6 +100,11 @@ struct llama_hparams { uint32_t n_group_used = 0; uint32_t n_group_experts = 0; + // MLA + SWA (i.e. dots3note) + uint32_t n_lora_kv_swa = 0; + uint32_t n_embd_head_k_mla_swa = 0; + uint32_t n_embd_head_v_mla_swa = 0; + float expert_group_scale = 0.05f; float expert_weights_scale = 0.0f; bool expert_weights_norm = false; diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 96cb045d2e..5192d6197d 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -20,6 +20,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( bool unified, uint32_t kv_size, uint32_t n_seq_max, + uint32_t n_ubatch, uint32_t n_pad, uint32_t n_swa, llama_swa_type swa_type, @@ -28,12 +29,41 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( const layer_reuse_cb & reuse) : hparams_lid(model.hparams), n_stream(unified ? 1 : n_seq_max) { + const bool has_swa = swa_type != LLAMA_SWA_TYPE_NONE; + + const layer_filter_cb filter_mla_full = [&](int32_t il) { + if (filter_mla && !filter_mla(il)) { + return false; + } + + return !(has_swa && model.hparams.is_swa(il)); + }; + + const layer_filter_cb filter_mla_swa = [&](int32_t il) { + if (filter_mla && !filter_mla(il)) { + return false; + } + + return model.hparams.is_swa(il); + }; + LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size); kv_mla = std::make_unique( model, model.hparams, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, - n_swa, swa_type, nullptr, filter_mla, reuse, nullptr); + 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_mla_full, reuse, nullptr); + + if (has_swa) { + const uint32_t size_swa = GGML_PAD(std::min(kv_size, n_swa*(unified ? n_seq_max : 1) + n_ubatch), 256); + + LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); + + kv_mla_swa = std::make_unique( + model, model.hparams, type_k, type_v, + v_trans, offload, unified, size_swa, n_seq_max, n_pad, + n_swa, swa_type, nullptr, filter_mla_swa, reuse, nullptr); + } // we use llama_kv_cache for caching indexer keys // by hand-tweaking some hparams we fool it to create @@ -47,14 +77,18 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); + // the lightning indexer only exists on full-attention layers, so no SWA handling here kv_lid = std::make_unique( model, hparams_lid, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, - n_swa, swa_type, nullptr, filter_lid, reuse, nullptr); + 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, reuse, nullptr); } void llama_kv_cache_dsa::clear(bool data) { kv_mla->clear(data); + if (kv_mla_swa) { + kv_mla_swa->clear(data); + } kv_lid->clear(data); } @@ -62,6 +96,9 @@ bool llama_kv_cache_dsa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) bool res = true; res = res & kv_mla->seq_rm(seq_id, p0, p1); + if (kv_mla_swa) { + res = res & kv_mla_swa->seq_rm(seq_id, p0, p1); + } res = res & kv_lid->seq_rm(seq_id, p0, p1); return res; @@ -69,21 +106,33 @@ bool llama_kv_cache_dsa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) void llama_kv_cache_dsa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { kv_mla->seq_cp(seq_id_src, seq_id_dst, p0, p1); + if (kv_mla_swa) { + kv_mla_swa->seq_cp(seq_id_src, seq_id_dst, p0, p1); + } kv_lid->seq_cp(seq_id_src, seq_id_dst, p0, p1); } void llama_kv_cache_dsa::seq_keep(llama_seq_id seq_id) { kv_mla->seq_keep(seq_id); + if (kv_mla_swa) { + kv_mla_swa->seq_keep(seq_id); + } kv_lid->seq_keep(seq_id); } void llama_kv_cache_dsa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { kv_mla->seq_add(seq_id, p0, p1, shift); + if (kv_mla_swa) { + kv_mla_swa->seq_add(seq_id, p0, p1, shift); + } kv_lid->seq_add(seq_id, p0, p1, shift); } void llama_kv_cache_dsa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { kv_mla->seq_div(seq_id, p0, p1, d); + if (kv_mla_swa) { + kv_mla_swa->seq_div(seq_id, p0, p1, d); + } kv_lid->seq_div(seq_id, p0, p1, d); } @@ -97,6 +146,11 @@ llama_pos llama_kv_cache_dsa::seq_pos_max(llama_seq_id seq_id) const { std::map llama_kv_cache_dsa::memory_breakdown() const { std::map mb = kv_mla->memory_breakdown(); + if (kv_mla_swa) { + for (const auto & buft_size : kv_mla_swa->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } for (const auto & buft_size : kv_lid->memory_breakdown()) { mb[buft_size.first] += buft_size.second; } @@ -133,6 +187,14 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch( break; } + llama_kv_cache::slot_info_vec_t sinfos_mla_swa; + if (kv_mla_swa) { + sinfos_mla_swa = kv_mla_swa->prepare(ubatches); + if (sinfos_mla_swa.empty()) { + break; + } + } + auto sinfos_lid = kv_lid->prepare(ubatches); if (sinfos_lid.empty()) { break; @@ -141,7 +203,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch( assert(sinfos_mla.size() == sinfos_lid.size()); return std::make_unique( - this, std::move(sinfos_mla), std::move(sinfos_lid), std::move(ubatches)); + this, std::move(sinfos_mla), std::move(sinfos_mla_swa), std::move(sinfos_lid), std::move(ubatches)); } while (false); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); @@ -157,17 +219,24 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_update(llama_context * lctx, b bool llama_kv_cache_dsa::get_can_shift() const { return kv_mla->get_can_shift() && + (!kv_mla_swa || kv_mla_swa->get_can_shift()) && kv_lid->get_can_shift() && kv_mla->get_size() == kv_lid->get_size(); } void llama_kv_cache_dsa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { kv_mla->state_write(io, seq_id, flags); + if (kv_mla_swa) { + kv_mla_swa->state_write(io, seq_id, flags); + } kv_lid->state_write(io, seq_id, flags); } void llama_kv_cache_dsa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { kv_mla->state_read(io, seq_id, flags); + if (kv_mla_swa) { + kv_mla_swa->state_read(io, seq_id, flags); + } kv_lid->state_read(io, seq_id, flags); } @@ -175,6 +244,10 @@ llama_kv_cache * llama_kv_cache_dsa::get_mla() const { return kv_mla.get(); } +llama_kv_cache * llama_kv_cache_dsa::get_mla_swa() const { + return kv_mla_swa.get(); +} + llama_kv_cache * llama_kv_cache_dsa::get_lid() const { return kv_lid.get(); } @@ -188,8 +261,11 @@ llama_kv_cache_dsa_context::llama_kv_cache_dsa_context(llama_memory_status statu llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( llama_kv_cache_dsa * kv) : ctx_mla(kv->get_mla()->init_full()), + ctx_mla_swa(kv->get_mla_swa() ? kv->get_mla_swa()->init_full() : nullptr), ctx_lid(kv->get_lid()->init_full()), - status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) { + status(llama_memory_status_combine( + llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status()), + ctx_mla_swa ? ctx_mla_swa->get_status() : LLAMA_MEMORY_STATUS_NO_UPDATE)) { } llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( @@ -197,20 +273,27 @@ llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( llama_context * lctx, bool optimize) : ctx_mla(kv->get_mla()->init_update(lctx, optimize)), + ctx_mla_swa(kv->get_mla_swa() ? kv->get_mla_swa()->init_update(lctx, optimize) : nullptr), ctx_lid(kv->get_lid()->init_update(lctx, optimize)), - status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) { + status(llama_memory_status_combine( + llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status()), + ctx_mla_swa ? ctx_mla_swa->get_status() : LLAMA_MEMORY_STATUS_NO_UPDATE)) { } llama_kv_cache_dsa_context::llama_kv_cache_dsa_context( llama_kv_cache_dsa * kv, slot_info_vec_t sinfos_mla, + slot_info_vec_t sinfos_mla_swa, slot_info_vec_t sinfos_lid, std::vector ubatches) : ubatches(std::move(ubatches)), // note: here we copy the ubatches. not sure if this is ideal ctx_mla(new llama_kv_cache_context(kv->get_mla(), std::move(sinfos_mla), this->ubatches)), + ctx_mla_swa(kv->get_mla_swa() ? new llama_kv_cache_context(kv->get_mla_swa(), std::move(sinfos_mla_swa), this->ubatches) : nullptr), ctx_lid(new llama_kv_cache_context(kv->get_lid(), std::move(sinfos_lid), this->ubatches)), - status(llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status())) { + status(llama_memory_status_combine( + llama_memory_status_combine(ctx_mla->get_status(), ctx_lid->get_status()), + ctx_mla_swa ? ctx_mla_swa->get_status() : LLAMA_MEMORY_STATUS_NO_UPDATE)) { } llama_kv_cache_dsa_context:: ~llama_kv_cache_dsa_context() = default; @@ -219,6 +302,9 @@ bool llama_kv_cache_dsa_context::next() { assert(status == LLAMA_MEMORY_STATUS_SUCCESS); ctx_mla->next(); + if (ctx_mla_swa) { + ctx_mla_swa->next(); + } ctx_lid->next(); if (++i_next >= ubatches.size()) { @@ -234,6 +320,9 @@ bool llama_kv_cache_dsa_context::apply() { bool res = true; res = res & ctx_mla->apply(); + if (ctx_mla_swa) { + res = res & ctx_mla_swa->apply(); + } res = res & ctx_lid->apply(); return res; @@ -255,6 +344,12 @@ const llama_kv_cache_context * llama_kv_cache_dsa_context::get_mla() const { return static_cast(ctx_mla.get()); } +const llama_kv_cache_context * llama_kv_cache_dsa_context::get_mla_swa() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return static_cast(ctx_mla_swa.get()); +} + const llama_kv_cache_context * llama_kv_cache_dsa_context::get_lid() const { assert(status == LLAMA_MEMORY_STATUS_SUCCESS); diff --git a/src/llama-kv-cache-dsa.h b/src/llama-kv-cache-dsa.h index e74fc4d910..b7b403a763 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -11,6 +11,7 @@ // utilizes two instances of llama_kv_cache: // - the first instance is for caching key tensors of the model, // - the second instance is for caching lightning indexer key tensors +// when swa_type != NONE, a third instance holds the SWA layers (iswa-style, window-sized) class llama_kv_cache_dsa : public llama_memory_i { public: @@ -23,6 +24,7 @@ public: bool unified, uint32_t kv_size, uint32_t n_seq_max, + uint32_t n_ubatch, uint32_t n_pad, uint32_t n_swa, llama_swa_type swa_type, @@ -69,8 +71,9 @@ public: // llama_kv_cache_dsa specific API // - llama_kv_cache * get_mla() const; - llama_kv_cache * get_lid() const; + llama_kv_cache * get_mla() const; + llama_kv_cache * get_mla_swa() const; // null when the model has no SWA layers + llama_kv_cache * get_lid() const; private: // we keep indexer KV cache hparams instance here as llama_kv_cache stores only reference to it @@ -78,6 +81,7 @@ private: const uint32_t n_stream = 1; std::unique_ptr kv_mla; + std::unique_ptr kv_mla_swa; std::unique_ptr kv_lid; }; @@ -102,6 +106,7 @@ public: llama_kv_cache_dsa_context( llama_kv_cache_dsa * kv, slot_info_vec_t sinfos_base, + slot_info_vec_t sinfos_swa, slot_info_vec_t sinfos_ik, std::vector ubatches); @@ -121,8 +126,9 @@ public: // llama_kv_cache_dsa_context specific API // - const llama_kv_cache_context * get_mla() const; - const llama_kv_cache_context * get_lid() const; + const llama_kv_cache_context * get_mla() const; + const llama_kv_cache_context * get_mla_swa() const; // null when the model has no SWA layers + const llama_kv_cache_context * get_lid() const; private: //llama_kv_cache_dsa * kv; @@ -133,6 +139,7 @@ private: std::vector ubatches; const llama_memory_context_ptr ctx_mla; + const llama_memory_context_ptr ctx_mla_swa; const llama_memory_context_ptr ctx_lid; const llama_memory_status status; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 5382cd7266..2e2bd7dc6d 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -323,7 +323,8 @@ 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 || model.arch == LLM_ARCH_GLM_DSA) && + if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_DOTS3NOTE) && hparams.n_embd_head_k_full == hparams.indexer_head_size) { attn_rot_k = true; } diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index abca773a9a..5e7341a66d 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -30,6 +30,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: + case LLM_ARCH_DOTS3NOTE: // TODO: need to handle SWA pattern and MLA+SWA config return false; default: return true; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c810055050..f0b48e2c7e 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -194,6 +194,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_deepseek2ocr(params); case LLM_ARCH_DEEPSEEK32: return new llama_model_deepseek32(params); + case LLM_ARCH_DOTS3NOTE: + return new llama_model_dots3note(params); case LLM_ARCH_DEEPSEEK4: return new llama_model_deepseek4(params); case LLM_ARCH_GLM_DSA: @@ -1907,7 +1909,9 @@ void llama_model::print_info() const { LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale); } - if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK2OCR || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4) { + if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK2OCR || + arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || + arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_MISTRAL4) { LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); LLAMA_LOG_INFO("%s: n_lora_q = %d\n", __func__, hparams.n_lora_q); LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv); @@ -2125,6 +2129,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_GLM_DSA: case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DOTS3NOTE: { if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) { // The NextN/MTP draft head runs dense MLA (no DSA indexer), so the @@ -2157,7 +2162,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (hparams.n_layer_nextn > 0) { filter_mla = [&](uint32_t il) { return il < hparams.n_layer(); }; } - llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (arch != LLM_ARCH_GLM_DSA || hparams.is_indexer_full(il)); }; + const bool lid_by_types = arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE; + llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (!lid_by_types || hparams.is_indexer_full(il)); }; res = new llama_kv_cache_dsa( *this, @@ -2168,6 +2174,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.kv_unified, cparams.n_ctx_seq, cparams.n_seq_max, + cparams.n_ubatch, 1, hparams.n_swa, hparams.swa_type, @@ -2641,6 +2648,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_LLAMA_EMBED: case LLM_ARCH_MAINCODER: case LLM_ARCH_GLM_DSA: + case LLM_ARCH_DOTS3NOTE: case LLM_ARCH_NANBEIGE: case LLM_ARCH_POCKETTTS: return LLAMA_ROPE_TYPE_NORM; diff --git a/src/models/dots3note.cpp b/src/models/dots3note.cpp new file mode 100644 index 0000000000..ba376dcc84 --- /dev/null +++ b/src/models/dots3note.cpp @@ -0,0 +1,510 @@ +#include "models.h" + +#include "llama-kv-cache.h" +#include "llama-kv-cache-dsa.h" + +// note: code adapted from deepseek32.cpp (DSA indexer + absorbed MLA) and step35.cpp (head-wise output gate) + +void llama_model_dots3note::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + hparams.f_norm_eps = 1e-6; // eps for the indexer k_norm layer norm + + // TODO: use MTP layer + 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_all"); + + // MoE parameters + ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert); + ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + 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); + + // MLA parameters of the full-attention layers + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + 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); + + // MLA parameters of the sliding-window layers + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK_SWA, hparams.n_lora_kv_swa); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA, hparams.n_embd_head_k_mla_swa); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA, hparams.n_embd_head_v_mla_swa); + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa); + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + + // DSA parameters - by default the indexer exists exactly on the full-attention layers + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_indexer_full_impl[il] = hparams.is_swa(il) ? 0 : 1; + } + ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); + + switch (hparams.n_layer()) { + case 46: type = LLM_TYPE_UNKNOWN; break; // 288B-A19B + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_dots3note::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + GGML_UNUSED(ml); + + if (!hparams.is_mla()) { + throw std::runtime_error("DOTS3NOTE architecture requires MLA"); + } + + const int64_t n_embd_head_qk_rope = hparams.n_rot(); + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = hparams.n_expert_shared; + + 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) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + const bool is_mtp = i >= n_layer; + // the NextN/MTP block uses the sliding-attention geometry + const bool is_swa = is_mtp || hparams.is_swa(i); + + // MTP tensors are preserved in the GGUF but there is no MTP graph yet + const int flags = is_mtp ? TENSOR_SKIP | TENSOR_NOT_REQUIRED : 0; + + const int64_t n_head_l = hparams.n_head(i); + + const int64_t kv_lora_rank = is_swa ? hparams.n_lora_kv_swa : hparams.n_lora_kv; + const int64_t n_embd_head_k_mla = is_swa ? hparams.n_embd_head_k_mla_swa : hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v_mla = is_swa ? hparams.n_embd_head_v_mla_swa : hparams.n_embd_head_v_mla(); + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags); + // norm applied on the shared rope key before rope + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_qk_rope}, flags); + + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head_l * n_embd_head_k_mla}, flags); + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, flags); + + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head_l}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head_l}, flags); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head_l * n_embd_head_v_mla, n_embd}, flags); + + // head-wise sigmoid output gate + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_l}, flags); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + + // DSA indexer + if (!is_mtp && hparams.is_indexer_full(i)) { + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, flags); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {hparams.indexer_head_size}, flags); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, hparams.indexer_head_size}, flags); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * hparams.indexer_head_size}, flags); + } + + if (is_mtp || i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags); + } else { + if (n_expert == 0 || n_expert_used == 0) { + throw std::runtime_error("n_expert and n_expert_used must be > 0"); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + } + + if (is_mtp) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags); + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, flags); + } + } +} + +std::unique_ptr llama_model_dots3note::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + GGML_ASSERT(hparams.is_mla()); + + const int64_t n_embd_head_qk_rope = hparams.n_rot(); + + 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; + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + 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(); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + + const bool is_swa = hparams.is_swa(il); + + const int64_t n_head_l = hparams.n_head(il); + + const int64_t kv_lora_rank = is_swa ? hparams.n_lora_kv_swa : hparams.n_lora_kv; + const int64_t n_embd_head_k_mla = is_swa ? hparams.n_embd_head_k_mla_swa : hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v_mla = is_swa ? hparams.n_embd_head_v_mla_swa : hparams.n_embd_head_v_mla(); + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + + const float kq_scale = 1.0f/sqrtf(float(n_embd_head_k_mla)); + const float freq_base_l = model.get_rope_freq_base(cparams, il); + + // norm + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self_attention + { + ggml_tensor * attn_inp = cur; + + 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 (full-attention layers only) + if (!is_swa) { + 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_NEOX, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(indexer_q_pe, "indexer_q_pe", il); + + 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_NEOX, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(indexer_k_pe, "indexer_k_pe", il); + + 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)); + + ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur); + cb(indexer_weights, "indexer_weights", il); + + 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 { + 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); + + indexer_score = ggml_relu(ctx0, indexer_kq); + cb(indexer_score, "indexer_score", il); + + 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); + + 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)); + 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_l, n_tokens} + ggml_tensor * q_nope = + ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head_l, n_tokens, ggml_row_size(q->type, n_embd_head_k_mla), + ggml_row_size(q->type, n_embd_head_k_mla) * n_head_l, 0); + cb(q_nope, "q_nope", il); + + // and {n_embd_head_qk_rope, n_head_l, n_tokens} + ggml_tensor * q_pe = ggml_view_3d( + ctx0, q, n_embd_head_qk_rope, n_head_l, n_tokens, ggml_row_size(q->type, n_embd_head_k_mla), + ggml_row_size(q->type, n_embd_head_k_mla) * n_head_l, 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); + + // norm on the shared rope key, applied before rope + k_pe = build_norm(k_pe, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + 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_l, 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_l, 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 with the absorption optimization + { + // {n_embd_head_qk_nope, n_tokens, n_head_l} + 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_l} x {n_embd_head_qk_nope, n_tokens, n_head_l} + 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_l, 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_l, n_tokens} + 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); + + // apply the head-wise output gate before o_proj, so wo stays out of build_attn + cur = build_attn(inp_attn_dsa, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il); + cb(cur, "attn_out", il); + + ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp); + cb(gate, "attn_gate", il); + + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "attn_gate_sigmoid", il); + + // broadcast the per-head gate over the head dimension + ggml_tensor * attn_3d = ggml_reshape_3d(ctx0, cur, n_embd_head_v_mla, n_head_l, n_tokens); + ggml_tensor * gate_3d = ggml_reshape_3d(ctx0, gate, 1, n_head_l, n_tokens); + attn_3d = ggml_mul(ctx0, attn_3d, gate_3d); + cb(attn_3d, "attn_gated", il); + + cur = ggml_reshape_2d(ctx0, attn_3d, n_embd_head_v_mla * n_head_l, n_tokens); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_output", 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 { + 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); + + 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); + + 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; + + cur = ggml_mul_mat(ctx0, model.output, cur); + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index ddb9ae2f12..e8247c53ff 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1156,6 +1156,18 @@ struct llama_model_deepseek32 : public llama_model_base { }; +struct llama_model_dots3note : public llama_model_base { + llama_model_dots3note(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 build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_deepseek4 : public llama_model_base { llama_model_deepseek4(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index e900bdc0da..b75740896b 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -104,6 +104,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } else if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA + || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_MISTRAL4) { n_embd = 128; @@ -163,6 +164,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA + || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_MISTRAL4) { ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576)); @@ -170,6 +172,15 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); + if (arch == LLM_ARCH_DOTS3NOTE) { + // SWA layers reuse the same MLA geometry as the full layers in this fixture + ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK_SWA, uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA, uint32_t(576)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA, uint32_t(192)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA, uint32_t(128)); + ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f); + } } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -192,7 +203,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f); // SWA pattern: every 5th layer is full attention (matches E2B layer_types) ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5)); - } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER) { + } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_DOTS3NOTE) { std::vector pattern; pattern.reserve(n_layer); for (uint32_t il = 0; il < n_layer; il++) { @@ -348,6 +359,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK: case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DOTS3NOTE: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: @@ -436,7 +448,7 @@ static bool arch_supported(const llm_arch arch) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE) { return false; } #endif // GGML_USE_WEBGPU