Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model: add dots3-note (#27060)

Browse Source 
* text: conversion

* init impl

* address review comments

* fix rope

* move to a new llama_kv_cache_dsa_iswa
This commit is contained in:
Xuan-Son Nguyen
2026-08-21 19:52:34 +02:00
committed by GitHub
parent 873e5d8e39
commit 5a32f7b66e
20 changed files with 1412 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-graph.cpp
+60 -6
View File
@@ -9,6 +9,7 @@
#include "llama-kv-cache.h"
#include "llama-kv-cache-iswa.h"
#include "llama-kv-cache-dsa.h"
#include "llama-kv-cache-dsa-iswa.h"
#include "llama-kv-cache-msa.h"
#include "llama-kv-cache-dsv4.h"
#include "llama-memory-hybrid.h"
@@ -507,10 +508,12 @@ void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) {
}
bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) {
const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);
mctx = static_cast<const llama_kv_cache_context *>(params.mctx);
this->mctx = mctx;
return can_reuse_impl(params);
}
bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) {
bool res = true;
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
@@ -567,10 +570,12 @@ void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
}
bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) {
const auto * mctx = static_cast<const llama_kv_cache_dsa_context *>(params.mctx);
mctx = static_cast<const llama_kv_cache_dsa_context *>(params.mctx);
this->mctx = mctx;
return can_reuse_impl(params);
}
bool llm_graph_input_attn_k_dsa::can_reuse_impl(const llm_graph_params & params) {
bool res = true;
res &= self_k_idxs_mla->ne[0] == params.ubatch.n_tokens;
@@ -582,6 +587,25 @@ bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) {
return res;
}
void llm_graph_input_attn_k_dsa_iswa::set_input(const llama_ubatch * ubatch) {
inp_dsa->set_input(ubatch);
inp_swa->set_input(ubatch);
}
bool llm_graph_input_attn_k_dsa_iswa::can_reuse(const llm_graph_params & params) {
mctx = static_cast<const llama_kv_cache_dsa_iswa_context *>(params.mctx);
inp_dsa->mctx = mctx->get_dsa();
inp_swa->mctx = mctx->get_swa();
bool res = true;
res &= inp_dsa->can_reuse_impl(params);
res &= inp_swa->can_reuse_impl(params);
return res;
}
void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
// base tensors may not be allocated if there are no non-SWA attention layers
if (self_k_idxs && self_k_idxs->buffer) {
@@ -3210,8 +3234,12 @@ ggml_tensor * llm_graph_context::build_attn(
return cur;
}
llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_context *>(mctx);
static std::unique_ptr<llm_graph_input_attn_k_dsa> build_attn_inp_k_dsa_impl(
ggml_context * ctx0,
const llama_ubatch & ubatch,
const llama_hparams & hparams,
const llama_cparams & cparams,
const llama_kv_cache_dsa_context * mctx_cur) {
auto inp = std::make_unique<llm_graph_input_attn_k_dsa>(hparams, cparams, mctx_cur);
@@ -3235,9 +3263,35 @@ llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
inp->self_k_rot_lid = mctx_cur->get_lid()->build_input_k_rot(ctx0);
}
return inp;
}
llm_graph_input_attn_k_dsa * llm_graph_context::build_attn_inp_k_dsa() const {
const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_context *>(mctx);
auto inp = build_attn_inp_k_dsa_impl(ctx0, ubatch, hparams, cparams, mctx_cur);
return (llm_graph_input_attn_k_dsa *) res->add_input(std::move(inp));
}
llm_graph_input_attn_k_dsa_iswa * llm_graph_context::build_attn_inp_k_dsa_iswa() const {
const auto * mctx_cur = static_cast<const llama_kv_cache_dsa_iswa_context *>(mctx);
auto inp_dsa = build_attn_inp_k_dsa_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_dsa());
// build_attn_inp_k_impl rejects SWA caches, so construct the input directly
auto inp_swa = std::make_unique<llm_graph_input_attn_k>(hparams, cparams, mctx_cur->get_swa());
inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
inp_swa->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
inp_swa->self_kq_mask_cnv = inp_swa->self_kq_mask;
auto inp = std::make_unique<llm_graph_input_attn_k_dsa_iswa>(std::move(inp_dsa), std::move(inp_swa), mctx_cur);
return (llm_graph_input_attn_k_dsa_iswa *) res->add_input(std::move(inp));
}
llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_enabled) const {
const auto * mctx_cur = static_cast<const llama_kv_cache_msa_context *>(mctx);