wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

Merge upstream master

Browse Source 
Assisted-by: Codex
This commit is contained in:
xrip committed 2026-10-06 12:03:55 +03:00
2 parents 74facb828e 1a3011cc0c
commit 411bf3c608
584 files changed
+69752 -32248
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/llama-graph.cpp
+191 -71
View File
@@ -75,13 +75,55 @@ void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens));
}
if (ubatch->embd) {
if (ubatch->embd && embd && !ubatch->is_mixed()) {
GGML_ASSERT(n_embd == embd->ne[0]);
const int64_t n_tokens = ubatch->n_tokens;
ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(embd));
}
if (ubatch->is_mixed() && embd) {
GGML_ASSERT(mixed_tokens && mixed_slots && mixed_embd && "mixed token/embd ubatch is not supported here");
std::vector<int32_t> ids;
std::vector<int64_t> slots;
for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
if (!ubatch->type[i]) {
ids.push_back(ubatch->token[i]);
slots.push_back(i);
}
}
GGML_ASSERT((int64_t) ids.size() == mixed_tokens->ne[0]);
GGML_ASSERT(n_embd == mixed_embd->ne[0]);
ggml_backend_tensor_set(mixed_tokens, ids.data(), 0, ggml_nbytes(mixed_tokens));
ggml_backend_tensor_set(mixed_slots, slots.data(), 0, ggml_nbytes(mixed_slots));
ggml_backend_tensor_set(mixed_embd, ubatch->embd, 0, ggml_nbytes(mixed_embd));
}
if (scale_rows) {
const int64_t n_tokens = ubatch->n_tokens;
std::vector<float> data(n_tokens);
for (int64_t i = 0; i < n_tokens; ++i) {
const bool is_embd = !ubatch->token || (ubatch->is_mixed() && ubatch->type[i]);
data[i] = is_embd ? 1.0f : scale_tok;
}
ggml_backend_tensor_set(scale_rows, data.data(), 0, ggml_nbytes(scale_rows));
}
}
// number of token rows of the mixed path, a non-mixed ubatch is sized for the worst case
static int64_t llm_graph_n_tok_rows(const llama_ubatch & ubatch) {
if (!ubatch.is_mixed()) {
return ubatch.n_tokens;
}
int64_t n = 0;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
n += !ubatch.type[i];
}
return n;
}
bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {
@@ -89,11 +131,16 @@ bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {
res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens);
res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens);
res &= (!mixed_tokens) || mixed_tokens->ne[0] == llm_graph_n_tok_rows(params.ubatch);
res &= (!mixed_embd) || mixed_embd->ne[1] == params.ubatch.n_tokens;
res &= (!scale_rows) || scale_rows->ne[1] == params.ubatch.n_tokens;
return res;
}
void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) {
ASSERT_EMBD_OR_TOKEN(*ubatch);
const int64_t n_tokens = ubatch->n_tokens;
if (ubatch->token) {
@@ -130,21 +177,7 @@ void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) {
if (ubatch->pos && pos) {
const int64_t n_tokens = ubatch->n_tokens;
if (ubatch->token && n_pos_per_embd == 4) {
// in case we're using M-RoPE with text tokens, convert the 1D positions to 4D
// the 3 first dims are the same, and 4th dim is all 0
std::vector<llama_pos> pos_data(n_tokens*n_pos_per_embd);
// copy the first dimension
for (int i = 0; i < n_tokens; ++i) {
pos_data[ i] = ubatch->pos[i];
pos_data[ n_tokens + i] = ubatch->pos[i];
pos_data[2 * n_tokens + i] = ubatch->pos[i];
pos_data[3 * n_tokens + i] = 0; // 4th dim is 0
}
ggml_backend_tensor_set(pos, pos_data.data(), 0, pos_data.size()*ggml_element_size(pos));
} else {
ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos));
}
ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos));
}
}
@@ -299,7 +332,7 @@ void llm_graph_input_cls::set_input(const llama_ubatch * ubatch) {
const bool last = (
cparams.pooling_type == LLAMA_POOLING_TYPE_LAST ||
(cparams.pooling_type == LLAMA_POOLING_TYPE_RANK && (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_QWEN3VL)) // qwen3 reranking & embedding models use last token
(cparams.pooling_type == LLAMA_POOLING_TYPE_RANK && cparams.causal_attn)
);
for (int i = 0; i < n_tokens; ++i) {
@@ -356,8 +389,7 @@ bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {
res &= s_copy->ne[0] == mctx->get_n_rs();
res &= s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= s_copy_extra->ne[0] == mctx->get_n_rs() - params.ubatch.n_seqs;
res &= s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= head == mctx->get_head();
res &= rs_z == mctx->get_rs_z();
@@ -473,8 +505,13 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {
}
void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
mctx->set_input_k_idxs(self_k_idxs, ubatch);
mctx->set_input_v_idxs(self_v_idxs, ubatch);
// the idxs are left unallocated when the KV cache is bypassed during training
if (self_k_idxs && self_k_idxs->buffer) {
mctx->set_input_k_idxs(self_k_idxs, ubatch);
}
if (self_v_idxs && self_v_idxs->buffer) {
mctx->set_input_v_idxs(self_v_idxs, ubatch);
}
// the mask is left unallocated when the graph only stores K/V without attending
// (e.g. DFlash's KV-injection pass)
@@ -1132,8 +1169,7 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs;
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= inp_rs->head == mctx->get_recr()->get_head();
res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();
@@ -1175,8 +1211,7 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) {
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs;
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= inp_rs->head == mctx->get_recr()->get_head();
res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();
@@ -1263,8 +1298,7 @@ bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params)
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs;
res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs;
res &= inp_rs->head == mctx->get_recr()->get_head();
res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();
@@ -1494,6 +1528,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
loras (params.loras),
mctx (params.mctx),
cross (params.cross),
prec_policy (params.prec_policy),
samplers (params.samplers),
cb_func (params.cb),
res (params.res),
@@ -1522,6 +1557,14 @@ ggml_tensor * llm_graph_context::build_lora_mm(
ggml_tensor * w_s) const {
ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);
if (prec_policy) {
prec_policy->apply(res);
}
if (w->type == GGML_TYPE_NVFP4) {
ggml_prec_set_acc(res, GGML_PREC_BF16);
}
if (w_s) {
res = ggml_mul(ctx0, res, w_s);
}
@@ -1554,6 +1597,14 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
ggml_tensor * w_s) const {
ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);
if (prec_policy) {
prec_policy->apply(res);
}
if (w->type == GGML_TYPE_NVFP4) {
ggml_prec_set_acc(res, GGML_PREC_BF16);
}
if (w_s) {
const int64_t n_expert = w_s->ne[0];
const int64_t n_tokens = cur->ne[2];
@@ -1836,7 +1887,7 @@ ggml_tensor * llm_graph_context::build_ffn(
const float limit = hparams.swiglu_clamp_shexp[il];
constexpr float eps = 1e-6f;
if (limit > eps) {
if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit);
} else {
tmp = ggml_clamp(ctx0, tmp, -limit, limit);
@@ -2260,7 +2311,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
const float limit = hparams.swiglu_clamp_exp[il];
constexpr float eps = 1e-6f;
if (limit > eps) {
if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
cur = ggml_swiglu_clamp(ctx0, cur, up, limit);
} else {
up = ggml_clamp(ctx0, up, -limit, limit);
@@ -2431,7 +2482,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
}
// input embeddings with optional lora
ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float tok_scale) const {
const int64_t n_embd_inp = hparams.n_embd_inp();
const int64_t n_embd = hparams.n_embd;
@@ -2448,15 +2499,9 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
cb(inp->embd, "inp_embd", -1);
ggml_set_input(inp->embd);
// select one of the 2 inputs, based on the batch contents
// ref: https://github.com/ggml-org/llama.cpp/pull/18550
std::array<ggml_tensor *, 2> inps;
// token embeddings path (ubatch.token != nullptr)
{
auto & cur = inps[0];
cur = ggml_get_rows(ctx0, tok_embd, inp->tokens);
// token embeddings with lora and padding
auto build_tok = [&](ggml_tensor * ids) {
ggml_tensor * cur = ggml_get_rows(ctx0, tok_embd, ids);
// apply lora for embedding tokens if needed
for (const auto & lora : *loras) {
@@ -2470,7 +2515,7 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
ggml_tensor * inpL_delta = ggml_scale(ctx0, ggml_mul_mat(
ctx0, lw->b, // non-transposed lora_b
ggml_get_rows(ctx0, lw->a, inp->tokens)
ggml_get_rows(ctx0, lw->a, ids)
), scale);
cur = ggml_add(ctx0, cur, inpL_delta);
@@ -2479,19 +2524,48 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
if (n_embd_inp != n_embd) {
cur = ggml_pad(ctx0, cur, hparams.n_embd_inp() - n_embd, 0, 0, 0);
}
}
return cur;
};
// select one of the 3 inputs, based on the batch contents
// ref: https://github.com/ggml-org/llama.cpp/pull/18550
std::array<ggml_tensor *, 3> inps = {};
// token embeddings path (ubatch.token != nullptr)
inps[0] = build_tok(inp->tokens);
// vector embeddings path (ubatch.embd != nullptr)
{
auto & cur = inps[1];
inps[1] = inp->embd;
cur = inp->embd;
// mixed path (ubatch.is_mixed()): set_rows the token rows into a copy of the embd rows, with its own inputs as select branches must not share tensors
// TODO: use inp->tokens and inp->embd once ggml_build_forward_select allows it
const bool has_mixed = llm_arch_supports_mixed_batch(arch) && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT;
if (has_mixed) {
const int64_t n_tok_rows = llm_graph_n_tok_rows(ubatch);
inp->mixed_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tok_rows);
cb(inp->mixed_tokens, "inp_mixed_tokens", -1);
ggml_set_input(inp->mixed_tokens);
inp->mixed_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_tok_rows);
cb(inp->mixed_slots, "inp_mixed_slots", -1);
ggml_set_input(inp->mixed_slots);
inp->mixed_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens);
cb(inp->mixed_embd, "inp_mixed_embd", -1);
ggml_set_input(inp->mixed_embd);
// note: set_rows writes into its destination, so it gets a copy of the input
inps[2] = ggml_set_rows(ctx0, ggml_dup(ctx0, inp->mixed_embd), build_tok(inp->mixed_tokens), inp->mixed_slots);
}
assert(ggml_are_same_shape (inps[0], inps[1]));
assert(ggml_are_same_stride(inps[0], inps[1]));
ggml_tensor * cur = ggml_build_forward_select(gf, inps.data(), inps.size(), ubatch.token ? 0 : 1);
const int idx = ubatch.is_mixed() ? 2 : ubatch.token ? 0 : 1;
ggml_tensor * cur = ggml_build_forward_select(gf, inps.data(), has_mixed ? 3 : 2, idx);
if (n_embd_inp != n_embd) {
cur = ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[1], 0);
@@ -2499,16 +2573,31 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
res->t_inp_embd = cur;
// For Granite architecture
// NOTE: For deepstack models, only apply scale to token inputs (ie text-only input).
// Raw embeddings are assumed to be multimodal inputs that should not be scaled.
if (hparams.f_embedding_scale != 0.0f && (ubatch.token || hparams.n_deepstack_layers == 0)) {
const bool scale_tok_only = hparams.f_embedding_scale != 0.0f && hparams.n_deepstack_layers > 0;
// For Granite architecture
if (hparams.f_embedding_scale != 0.0f && !scale_tok_only) {
if (!ggml_is_contiguous(cur)) {
cur = ggml_cont(ctx0, cur);
}
cur = ggml_scale(ctx0, cur, hparams.f_embedding_scale);
}
// scale the token rows only, applied after the select so that the graph is the same for any batch contents
inp->scale_tok = tok_scale*(scale_tok_only ? hparams.f_embedding_scale : 1.0f);
if (inp->scale_tok != 1.0f) {
inp->scale_rows = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, ubatch.n_tokens);
cb(inp->scale_rows, "inp_scale_rows", -1);
ggml_set_input(inp->scale_rows);
if (!ggml_is_contiguous(cur)) {
cur = ggml_cont(ctx0, cur);
}
cur = ggml_mul(ctx0, cur, inp->scale_rows);
}
cb(cur, "embd", -1);
res->add_input(std::move(inp));
@@ -2687,7 +2776,10 @@ ggml_tensor * llm_graph_context::build_attn_mha(
// split the batch into streams if needed
const auto n_stream = k->ne[3];
q = ggml_view_4d(ctx0, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, q->nb[1], q->nb[2], q->nb[3]/n_stream, 0);
// the stream dim steps over one stream's worth of dim-2 (tokens): (ne[2]/n_stream) rows of stride nb[2].
// q->nb[3]/n_stream only equals that for a contiguous q; nope-only MLA (glm5-next) passes a permuted
// q_absorbed where nb[3] != ne[2]*nb[2], so using nb[3] read another head's queries for streams s >= 1.
q = ggml_view_4d(ctx0, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, q->nb[1], q->nb[2], q->nb[2]*(q->ne[2]/n_stream), 0);
q = ggml_permute(ctx0, q, 0, 2, 1, 3);
k = ggml_permute(ctx0, k, 0, 2, 1, 3);
@@ -2953,21 +3045,30 @@ ggml_tensor * llm_graph_context::build_attn(
const auto * mctx_cur = inp->mctx;
// store to KV cache
{
const auto & k_idxs = inp->get_k_idxs();
const auto & v_idxs = inp->get_v_idxs();
ggml_tensor * q = q_cur;
ggml_tensor * k;
ggml_tensor * v;
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
if (cparams.training) {
GGML_ASSERT(mctx_cur->get_n_kv() == n_tokens);
k = k_cur;
v = v_cur;
} else {
{
const auto & k_idxs = inp->get_k_idxs();
const auto & v_idxs = inp->get_v_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}
k = mctx_cur->get_k(ctx0, il);
v = mctx_cur->get_v(ctx0, il);
}
ggml_tensor * kq_mask = inp->get_kq_mask();
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
@@ -3204,14 +3305,22 @@ ggml_tensor * llm_graph_context::build_attn(
const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
// whole seq fits into batch in training mode
const bool use_kv_cur = cparams.training && k_cur && v_cur;
if (use_kv_cur) {
GGML_ASSERT(mctx_cur->get_n_kv() == n_tokens);
}
const bool store_kv = !use_kv_cur || hparams.n_layer_kv_from_start >= 0;
// optionally store to KV cache
if (k_cur) {
if (store_kv && k_cur) {
const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
}
if (v_cur) {
if (store_kv && v_cur) {
const auto & v_idxs = is_swa ? inp->get_v_idxs_swa() : inp->get_v_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
@@ -3220,8 +3329,8 @@ ggml_tensor * llm_graph_context::build_attn(
const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = mctx_cur->get_v(ctx0, il);
ggml_tensor * k = use_kv_cur ? k_cur : mctx_cur->get_k(ctx0, il);
ggml_tensor * v = use_kv_cur ? v_cur : mctx_cur->get_v(ctx0, il);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
@@ -3549,8 +3658,8 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
ggml_tensor * llm_graph_context::build_rs(
ggml_tensor * s,
ggml_tensor * state_copy,
ggml_tensor * state_copy_main,
ggml_tensor * state_copy_extra,
int32_t state_size,
int32_t n_seqs,
uint32_t n_rs,
@@ -3569,12 +3678,19 @@ ggml_tensor * llm_graph_context::build_rs(
// copy states
// NOTE: assuming the copy destinations are ALL contained between rs_head and rs_head + n_rs
// {state_size, rs_size} -> {state_size, n_seqs}
ggml_tensor * output_states = get_state_rows(ctx0, states, state_copy_main);
// one gather of the states i0..n_rs (ubatch states then extra states), sized by n_rs so the reserve covers every split
// {state_size, rs_size} -> {state_size, n_rs - i0}
const int64_t i0 = n_rs - state_copy->ne[0];
ggml_tensor * states_all = ggml_get_rows(ctx0, states, state_copy);
ggml_tensor * output_states = get_state_rows ?
get_state_rows(ctx0, states, state_copy_main) :
ggml_view_2d(ctx0, states_all, state_size, n_seqs, states_all->nb[1], 0);
ggml_build_forward_expand(gf, output_states);
// copy extra states which won't be changed further (between n_seqs and n_rs)
ggml_tensor * states_extra = ggml_get_rows(ctx0, states, state_copy_extra);
ggml_tensor * states_extra = ggml_view_2d(ctx0, states_all, state_size, n_rs - n_seqs, states_all->nb[1], (n_seqs - i0)*states_all->nb[1]);
ggml_build_forward_expand(gf,
ggml_cpy(ctx0,
states_extra,
@@ -3601,8 +3717,8 @@ static std::unique_ptr<llm_graph_input_rs> build_rs_inp_impl(
ggml_set_input(inp->txn_copy);
}
inp->s_copy_main = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0);
inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]);
inp->s_copy_main = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0);
inp->s_copy_tail = ggml_view_1d(ctx0, inp->s_copy, n_rs - 1, inp->s_copy->nb[0]);
inp->head = mctx_cur->get_head();
inp->rs_z = mctx_cur->get_rs_z();
@@ -3626,7 +3742,11 @@ ggml_tensor * llm_graph_context::build_rs(
const llm_graph_get_rows_fn & get_state_rows) const {
const auto * kv_state = inp->mctx;
return build_rs(s, inp->s_copy_main, inp->s_copy_extra, state_size, n_seqs,
// a custom getter reads the states of the ubatch straight from the cache, so the gather skips the first
// state: it still holds the n_rs - n_seqs extra states and copies no state of a single sequence ubatch
ggml_tensor * state_copy = get_state_rows ? inp->s_copy_tail : inp->s_copy;
return build_rs(s, state_copy, inp->s_copy_main, state_size, n_seqs,
kv_state->get_n_rs(), kv_state->get_head(), kv_state->get_size(), kv_state->get_rs_z(),
get_state_rows);
}
@@ -3794,8 +3914,8 @@ void llm_graph_context::build_pooling(
} break;
case LLAMA_POOLING_TYPE_RANK:
{
if (arch == LLM_ARCH_MODERN_BERT) {
// modern bert gte reranker builds mean first then applies prediction head and classifier
if (hparams.pooling_type_cls == LLAMA_POOLING_TYPE_MEAN) {
// modern bert with classifier_pooling = "mean" builds mean first then applies prediction head and classifier
// https://github.com/huggingface/transformers/blob/main/src/transformers/models/modernbert/modular_modernbert.py#L1404-1411
ggml_tensor * inp_mean = build_inp_mean();
cur = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, inp)), inp_mean);