#include "llama-graph.h"

#include "llama-impl.h"

#include "llama-batch.h"

#include "llama-cparams.h"

#include "llama-kv-cache.h"

#include "llama-kv-cache-iswa.h"

#include "llama-memory-hybrid.h"

#include "llama-memory-recurrent.h"

#include <cassert>

#include <cmath>

#include <cstring>

#include <unordered_set>

void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {

if (ubatch->token) {

const int64_t n_tokens = ubatch->n_tokens;

ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens));

}

if (ubatch->embd) {

const int64_t 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));

}

}

bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {

bool res = true;

res &= (!tokens && !params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens);

res &= (!embd && !params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens);

return res;

}

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));

}

}

}

bool llm_graph_input_pos::can_reuse(const llm_graph_params & params) {

bool res = true;

res &= pos->ne[0] == params.ubatch.n_tokens*n_pos_per_embd;

return res;

}

void llm_graph_input_attn_temp::set_input(const llama_ubatch * ubatch) {

if (ubatch->pos && attn_scale) {

const int64_t n_tokens = ubatch->n_tokens;

GGML_ASSERT(f_attn_temp_scale != 0.0f);

GGML_ASSERT(n_attn_temp_floor_scale != 0);

std::vector<float> attn_scale_data(n_tokens, 0.0f);

for (int i = 0; i < n_tokens; ++i) {

const float pos = ubatch->pos[i];

attn_scale_data[i] = std::log(

std::floor((pos + f_attn_temp_offset) / n_attn_temp_floor_scale) + 1.0

) * f_attn_temp_scale + 1.0;

}

ggml_backend_tensor_set(attn_scale, attn_scale_data.data(), 0, n_tokens*ggml_element_size(attn_scale));

}

}

void llm_graph_input_pos_bucket::set_input(const llama_ubatch * ubatch) {

if (pos_bucket) {

const int64_t n_tokens = ubatch->n_tokens;

GGML_ASSERT(ggml_backend_buffer_is_host(pos_bucket->buffer));

GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing

int32_t * data = (int32_t *) pos_bucket->data;

for (int j = 0; j < n_tokens; ++j) {

for (int i = 0; i < n_tokens; ++i) {

data[j*n_tokens + i] = llama_relative_position_bucket(ubatch->pos[i], ubatch->pos[j], hparams.n_rel_attn_bkts, true);

}

}

}

}

void llm_graph_input_pos_bucket_kv::set_input(const llama_ubatch * ubatch) {

if (pos_bucket) {

mctx->set_input_pos_bucket(pos_bucket, ubatch);

}

}

void llm_graph_input_out_ids::set_input(const llama_ubatch * ubatch) {

GGML_ASSERT(out_ids);

const int64_t n_tokens = ubatch->n_tokens;

GGML_ASSERT(ggml_backend_buffer_is_host(out_ids->buffer));

int32_t * data = (int32_t *) out_ids->data;

if (n_outputs == n_tokens) {

for (int i = 0; i < n_tokens; ++i) {

data[i] = i;

}

return;

}

GGML_ASSERT(ubatch->output);

int n_outputs = 0;

for (int i = 0; i < n_tokens; ++i) {

if (ubatch->output[i]) {

data[n_outputs++] = i;

}

}

}

bool llm_graph_input_out_ids::can_reuse(const llm_graph_params & params) {

bool res = true;

res &= n_outputs == params.n_outputs;

return res;

}

void llm_graph_input_mean::set_input(const llama_ubatch * ubatch) {

if (cparams.embeddings && cparams.pooling_type == LLAMA_POOLING_TYPE_MEAN) {

const int64_t n_tokens = ubatch->n_tokens;

const int64_t n_seq_tokens = ubatch->n_seq_tokens;

const int64_t n_seqs_unq = ubatch->n_seqs_unq;

GGML_ASSERT(mean);

GGML_ASSERT(ggml_backend_buffer_is_host(mean->buffer));

float * data = (float *) mean->data;

memset(mean->data, 0, n_tokens*n_seqs_unq*ggml_element_size(mean));

std::vector<uint64_t> sums(n_seqs_unq, 0);

for (int i = 0; i < n_tokens; i += n_seq_tokens) {

for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {

const llama_seq_id seq_id = ubatch->seq_id[i][s];

const int32_t seq_idx = ubatch->seq_idx[seq_id];

sums[seq_idx] += ubatch->n_seq_tokens;

}

}

std::vector<float> div(n_seqs_unq, 0.0f);

for (int s = 0; s < n_seqs_unq; ++s) {

const uint64_t sum = sums[s];

if (sum > 0) {

div[s] = 1.0f/float(sum);

}

}

for (int i = 0; i < n_tokens; i += n_seq_tokens) {

for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {

const llama_seq_id seq_id = ubatch->seq_id[i][s];

const int32_t seq_idx = ubatch->seq_idx[seq_id];

for (int j = 0; j < n_seq_tokens; ++j) {

data[seq_idx*n_tokens + i + j] = div[seq_idx];

}

}

}

}

}

void llm_graph_input_cls::set_input(const llama_ubatch * ubatch) {

const int64_t n_tokens = ubatch->n_tokens;

const int64_t n_seqs_unq = ubatch->n_seqs_unq;

if (cparams.embeddings && (

cparams.pooling_type == LLAMA_POOLING_TYPE_CLS ||

cparams.pooling_type == LLAMA_POOLING_TYPE_RANK ||

cparams.pooling_type == LLAMA_POOLING_TYPE_LAST

)) {

GGML_ASSERT(cls);

GGML_ASSERT(ggml_backend_buffer_is_host(cls->buffer));

uint32_t * data = (uint32_t *) cls->data;

memset(cls->data, 0, n_seqs_unq*ggml_element_size(cls));

std::vector<int> target_pos(n_seqs_unq, -1);

std::vector<int> target_row(n_seqs_unq, -1);

const bool last = (

cparams.pooling_type == LLAMA_POOLING_TYPE_LAST ||

(cparams.pooling_type == LLAMA_POOLING_TYPE_RANK && arch == LLM_ARCH_QWEN3) // qwen3 reranking & embedding models use last token

);

for (int i = 0; i < n_tokens; ++i) {

const llama_pos pos = ubatch->pos[i];

for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {

const llama_seq_id seq_id = ubatch->seq_id[i][s];

const int32_t seq_idx = ubatch->seq_idx[seq_id];

if (

(target_pos[seq_idx] == -1) ||

( last && pos >= target_pos[seq_idx]) ||

(!last && pos < target_pos[seq_idx])

) {

target_pos[seq_idx] = pos;

target_row[seq_idx] = i;

}

}

}

for (int s = 0; s < n_seqs_unq; ++s) {

if (target_row[s] >= 0) {

data[s] = target_row[s];

}

}

}

}

void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {

GGML_UNUSED(ubatch);

const int64_t n_rs = mctx->get_n_rs();

if (s_copy) {

GGML_ASSERT(ggml_backend_buffer_is_host(s_copy->buffer));

int32_t * data = (int32_t *) s_copy->data;

// assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n

for (uint32_t i = 0; i < n_rs; ++i) {

data[i] = mctx->s_copy(i);

}

}

}

bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {

const auto * mctx = static_cast<const llama_memory_recurrent_context *>(params.mctx);

this->mctx = mctx;

bool res = true;

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 &= head == mctx->get_head();

res &= rs_z == mctx->get_rs_z();

return res;

}

void llm_graph_input_cross_embd::set_input(const llama_ubatch * ubatch) {

GGML_UNUSED(ubatch);

if (cross_embd && !cross->v_embd.empty()) {

assert(cross_embd->type == GGML_TYPE_F32);

ggml_backend_tensor_set(cross_embd, cross->v_embd.data(), 0, ggml_nbytes(cross_embd));

}

}

static void print_mask(const float * data, int64_t n_tokens, int64_t n_kv, int64_t n_swa, llama_swa_type swa_type) {

LLAMA_LOG_DEBUG("%s: === Attention mask ===\n", __func__);

const char * swa_type_str = "unknown";

switch (swa_type) {

case LLAMA_SWA_TYPE_NONE: swa_type_str = "LLAMA_SWA_TYPE_NONE"; break;

case LLAMA_SWA_TYPE_STANDARD: swa_type_str = "LLAMA_SWA_TYPE_STANDARD"; break;

case LLAMA_SWA_TYPE_CHUNKED: swa_type_str = "LLAMA_SWA_TYPE_CHUNKED"; break;

case LLAMA_SWA_TYPE_SYMMETRIC: swa_type_str = "LLAMA_SWA_TYPE_SYMMETRIC"; break;

};

LLAMA_LOG_DEBUG("%s: n_swa : %d, n_kv: %d, swq_type: %s\n", __func__, (int)n_swa, (int)n_kv, swa_type_str);

LLAMA_LOG_DEBUG("%s: '0' = can attend, '∞' = masked\n", __func__);

LLAMA_LOG_DEBUG("%s: Rows = query tokens, Columns = key/value tokens\n\n", __func__);

LLAMA_LOG_DEBUG(" ");

for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) {

LLAMA_LOG_DEBUG("%2d", j);

}

LLAMA_LOG_DEBUG("\n");

for (int i = 0; i < std::min((int64_t)20, n_tokens); ++i) {

LLAMA_LOG_DEBUG(" %2d ", i);

for (int j = 0; j < std::min((int64_t)20, n_kv); ++j) {

float val = data[i * n_kv + j];

if (val == -INFINITY) {

LLAMA_LOG_DEBUG(" ∞");

} else {

LLAMA_LOG_DEBUG(" 0");

}

}

LLAMA_LOG_DEBUG("\n");

}

}

void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {

const int64_t n_kv = ubatch->n_tokens;

const int64_t n_tokens = ubatch->n_tokens;

const auto fill_mask = [&](float * data, int n_swa, llama_swa_type swa_type) {

for (int i1 = 0; i1 < n_tokens; ++i1) {

const llama_seq_id s1 = ubatch->seq_id[i1][0];

const llama_pos p1 = ubatch->pos[i1];

const uint64_t idst = i1*n_kv;

for (int i0 = 0; i0 < n_tokens; ++i0) {

const llama_seq_id s0 = ubatch->seq_id[i0][0];

const llama_pos p0 = ubatch->pos[i0];

// mask different sequences

if (s0 != s1) {

continue;

}

// mask future tokens

if (cparams.causal_attn && p0 > p1) {

continue;

}

// apply SWA if any

if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {

continue;

}

data[idst + i0] = hparams.use_alibi ? -std::abs(p0 - p1) : 0.0f;

}

}

};

{

GGML_ASSERT(self_kq_mask);

GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));

float * data = (float *) self_kq_mask->data;

std::fill(data, data + ggml_nelements(self_kq_mask), -INFINITY);

fill_mask(data, 0, LLAMA_SWA_TYPE_NONE);

if (debug) {

print_mask(data, n_tokens, n_kv, 0, LLAMA_SWA_TYPE_NONE);

}

}

if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {

GGML_ASSERT(self_kq_mask_swa);

GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer));

float * data = (float *) self_kq_mask_swa->data;

std::fill(data, data + ggml_nelements(self_kq_mask_swa), -INFINITY);

fill_mask(data, hparams.n_swa, hparams.swa_type);

if (debug) {

print_mask(data, n_tokens, n_kv, hparams.n_swa, hparams.swa_type);

}

}

}

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);

mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);

}

bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) {

const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);

this->mctx = mctx;

bool res = true;

res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;

//res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there

res &= self_kq_mask->ne[0] == mctx->get_n_kv();

res &= self_kq_mask->ne[1] == params.ubatch.n_tokens;

return res;

}

void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {

mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch);

mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch);

mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);

mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch);

mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch);

mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);

}

bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {

const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);

this->mctx = mctx;

bool res = true;

res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;

//res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there

res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;

//res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there

res &= self_kq_mask->ne[0] == mctx->get_base()->get_n_kv();

res &= self_kq_mask->ne[1] == params.ubatch.n_tokens;

res &= self_kq_mask_swa->ne[0] == mctx->get_swa()->get_n_kv();

res &= self_kq_mask_swa->ne[1] == params.ubatch.n_tokens;

return res;

}

void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) {

GGML_ASSERT(cross_kq_mask);

const int64_t n_enc = cross_kq_mask->ne[0];

const int64_t n_tokens = ubatch->n_tokens;

GGML_ASSERT(ggml_backend_buffer_is_host(cross_kq_mask->buffer));

GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing

float * data = (float *) cross_kq_mask->data;

for (int i = 0; i < n_tokens; ++i) {

for (int j = 0; j < n_enc; ++j) {

float f = -INFINITY;

for (int s = 0; s < ubatch->n_seq_id[i]; ++s) {

const llama_seq_id seq_id = ubatch->seq_id[i][s];

if (cross->seq_ids_enc[j].find(seq_id) != cross->seq_ids_enc[j].end()) {

f = 0.0f;

}

}

data[i*n_enc + j] = f;

}

}

}

void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {

mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);

mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);

mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);

const int64_t n_rs = mctx->get_recr()->get_n_rs();

if (inp_rs->s_copy) {

GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer));

int32_t * data = (int32_t *) inp_rs->s_copy->data;

// assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n

for (uint32_t i = 0; i < n_rs; ++i) {

data[i] = mctx->get_recr()->s_copy(i);

}

}

}

bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {

const auto * mctx = static_cast<const llama_memory_hybrid_context *>(params.mctx);

this->mctx = mctx;

bool res = true;

res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;

//res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there

res &= inp_attn->self_kq_mask->ne[0] == mctx->get_attn()->get_n_kv();

res &= inp_attn->self_kq_mask->ne[1] == params.ubatch.n_tokens;

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->head == mctx->get_recr()->get_head();

res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();

return res;

}

void llm_graph_input_sampling::set_input(const llama_ubatch * ubatch) {

// set the inputs only for the active samplers in the current ubatch

std::unordered_set<llama_seq_id> active_samplers;

for (uint32_t i = 0; i < ubatch->n_tokens; i++) {

if (ubatch->output[i]) {

llama_seq_id seq_id = ubatch->seq_id[i][0];

active_samplers.insert(seq_id);

}

}

for (auto seq_id : active_samplers) {

if (samplers.find(seq_id) == samplers.end()) {

continue;

}

auto & sampler = samplers[seq_id];

if (sampler->iface->backend_set_input) {

sampler->iface->backend_set_input(sampler);

}

}

}

bool llm_graph_input_sampling::can_reuse(const llm_graph_params & params) {

if (samplers.size() != params.samplers.size()) {

return false;

}

for (const auto & [seq_id, sampler] : params.samplers) {

if (samplers[seq_id] != sampler) {

return false;

}

}

return true;

}

//

// llm_graph_result

//

llm_graph_result::llm_graph_result(int64_t max_nodes) : max_nodes(max_nodes) {

reset();

const char * LLAMA_GRAPH_RESULT_DEBUG = getenv("LLAMA_GRAPH_RESULT_DEBUG");

debug = LLAMA_GRAPH_RESULT_DEBUG ? atoi(LLAMA_GRAPH_RESULT_DEBUG) : 0;

}

int64_t llm_graph_result::get_max_nodes() const {

return max_nodes;

}

void llm_graph_result::reset() {

t_tokens = nullptr;

t_logits = nullptr;

t_embd = nullptr;

t_embd_pooled = nullptr;

t_sampled.clear();

t_sampled_probs.clear();

t_sampled_logits.clear();

t_candidates.clear();

params = {};

inputs.clear();

buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));

ggml_init_params params = {

/*.mem_size =*/ buf_compute_meta.size(),

/*.mem_buffer =*/ buf_compute_meta.data(),

/*.no_alloc =*/ true,

};

ctx_compute.reset(ggml_init(params));

gf = ggml_new_graph_custom(ctx_compute.get(), max_nodes, false);

}

void llm_graph_result::set_inputs(const llama_ubatch * ubatch) {

for (auto & input : inputs) {

input->set_input(ubatch);

}

}

void llm_graph_result::set_outputs() {

if (t_logits != nullptr) {

ggml_set_output(t_logits);

}

if (t_embd != nullptr) {

ggml_set_output(t_embd);

}

if (t_embd_pooled != nullptr) {

ggml_set_output(t_embd_pooled);

}

for (auto & [seq_id, t] : t_sampled) {

if (t != nullptr) {

ggml_set_output(t);

}

}

for (auto & [seq_id, t] : t_sampled_probs) {

if (t != nullptr) {

ggml_set_output(t);

}

}

for (auto & [seq_id, t] : t_sampled_logits) {

if (t != nullptr) {

ggml_set_output(t);

}

}

for (auto & [seq_id, t] : t_candidates) {

if (t != nullptr) {

ggml_set_output(t);

}

}

}

bool llm_graph_result::can_reuse(const llm_graph_params & params) {

if (!this->params.allow_reuse(params)) {

if (debug > 1) {

LLAMA_LOG_DEBUG("%s: cannot reuse graph due to incompatible graph parameters\n", __func__);

}

return false;

}

if (debug > 1) {

LLAMA_LOG_DEBUG("%s: checking compatibility of %d inputs:\n", __func__, (int) inputs.size());

}

bool res = true;

for (auto & input : inputs) {

const bool cur = input->can_reuse(params);

if (debug > 1) {

LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", "placeholder", cur);

}

res = res && cur;

}

if (debug > 0) {

LLAMA_LOG_DEBUG("%s: can reuse graph = %d\n", __func__, res);

}

return res;

}

llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) {

inputs.emplace_back(std::move(input));

return inputs.back().get();

}

void llm_graph_result::set_params(const llm_graph_params & params) {

this->params = params;

}

//

// llm_graph_context

//

llm_graph_context::llm_graph_context(const llm_graph_params & params) :

arch (params.arch),

hparams (params.hparams),

cparams (params.cparams),

ubatch (params.ubatch),

n_embd (hparams.n_embd),

n_layer (hparams.n_layer),

n_rot (hparams.n_rot),

n_ctx (cparams.n_ctx),

n_head (hparams.n_head()),

n_head_kv (hparams.n_head_kv()),

n_embd_head_k (hparams.n_embd_head_k),

n_embd_k_gqa (hparams.n_embd_k_gqa()),

n_embd_head_v (hparams.n_embd_head_v),

n_embd_v_gqa (hparams.n_embd_v_gqa()),

n_expert (hparams.n_expert),

n_expert_used (cparams.warmup ? hparams.n_expert : hparams.n_expert_used),

freq_base (cparams.rope_freq_base),

freq_scale (cparams.rope_freq_scale),

ext_factor (cparams.yarn_ext_factor),

attn_factor (cparams.yarn_attn_factor),

beta_fast (cparams.yarn_beta_fast),

beta_slow (cparams.yarn_beta_slow),

norm_eps (hparams.f_norm_eps),

norm_rms_eps (hparams.f_norm_rms_eps),

n_tokens (ubatch.n_tokens),

n_outputs (params.n_outputs),

n_ctx_orig (cparams.n_ctx_orig_yarn),

pooling_type (cparams.pooling_type),

rope_type (hparams.rope_type),

sched (params.sched),

backend_cpu (params.backend_cpu),

cvec (params.cvec),

loras (params.loras),

mctx (params.mctx),

cross (params.cross),

samplers (params.samplers),

cb_func (params.cb),

res (params.res),

ctx0 (res->get_ctx()),

gf (res->get_gf()) {

res->set_params(params);

}

void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const {

if (cb_func) {

cb_func(ubatch, cur, name, il);

}

}

ggml_tensor * llm_graph_context::build_cvec(

ggml_tensor * cur,

int il) const {

return cvec->apply_to(ctx0, cur, il);

}

ggml_tensor * llm_graph_context::build_lora_mm(

ggml_tensor * w,

ggml_tensor * cur) const {

ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);

for (const auto & lora : *loras) {

llama_adapter_lora_weight * lw = lora.first->get_weight(w);

if (lw == nullptr) {

continue;

}

const float adapter_scale = lora.second;

const float scale = lw->get_scale(lora.first->alpha, adapter_scale);

ggml_tensor * ab_cur = ggml_mul_mat(

ctx0, lw->b,

ggml_mul_mat(ctx0, lw->a, cur)

);

ab_cur = ggml_scale(ctx0, ab_cur, scale);

res = ggml_add(ctx0, res, ab_cur);

}

return res;

}

ggml_tensor * llm_graph_context::build_lora_mm_id(

ggml_tensor * w, // ggml_tensor * as

ggml_tensor * cur, // ggml_tensor * b

ggml_tensor * ids) const {

ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);

for (const auto & lora : *loras) {

llama_adapter_lora_weight * lw = lora.first->get_weight(w);

if (lw == nullptr) {

continue;

}

const float alpha = lora.first->alpha;

const float rank = (float) lw->b->ne[0];

const float scale = alpha ? lora.second * alpha / rank : lora.second;

ggml_tensor * ab_cur = ggml_mul_mat_id(

ctx0, lw->b,

ggml_mul_mat_id(ctx0, lw->a, cur, ids),

ids

);

ab_cur = ggml_scale(ctx0, ab_cur, scale);

res = ggml_add(ctx0, res, ab_cur);

}

return res;

}

ggml_tensor * llm_graph_context::build_norm(

ggml_tensor * cur,

ggml_tensor * mw,

ggml_tensor * mb,

llm_norm_type type,

int il) const {

switch (type) {

case LLM_NORM: cur = ggml_norm (ctx0, cur, hparams.f_norm_eps); break;

case LLM_NORM_RMS: cur = ggml_rms_norm(ctx0, cur, hparams.f_norm_rms_eps); break;

case LLM_NORM_GROUP:

{

cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, cur->ne[1]);

cur = ggml_group_norm(ctx0, cur, hparams.n_norm_groups, hparams.f_norm_group_eps);

cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], cur->ne[2]);

} break;

}

if (mw || mb) {

cb(cur, "norm", il);

}

if (mw) {

cur = ggml_mul(ctx0, cur, mw);

if (mb) {

cb(cur, "norm_w", il);

}

}

if (mb) {

cur = ggml_add(ctx0, cur, mb);

}

return cur;

}

ggml_tensor * llm_graph_context::build_ffn(

ggml_tensor * cur,

ggml_tensor * up,

ggml_tensor * up_b,

ggml_tensor * up_s,

ggml_tensor * gate,

ggml_tensor * gate_b,

ggml_tensor * gate_s,

ggml_tensor * down,

ggml_tensor * down_b,

ggml_tensor * down_s,

ggml_tensor * act_scales,

llm_ffn_op_type type_op,

llm_ffn_gate_type type_gate,

int il) const {

ggml_tensor * tmp = up ? build_lora_mm(up, cur) : cur;

cb(tmp, "ffn_up", il);

if (up_b) {

tmp = ggml_add(ctx0, tmp, up_b);

cb(tmp, "ffn_up_b", il);

}

if (up_s) {

tmp = ggml_mul(ctx0, tmp, up_s);

cb(tmp, "ffn_up_s", il);

}

if (gate) {

switch (type_gate) {

case LLM_FFN_SEQ:

{

cur = build_lora_mm(gate, tmp);

cb(cur, "ffn_gate", il);

} break;

case LLM_FFN_PAR:

{

cur = build_lora_mm(gate, cur);

cb(cur, "ffn_gate", il);

} break;

}

if (gate_b) {

cur = ggml_add(ctx0, cur, gate_b);

cb(cur, "ffn_gate_b", il);

}

if (gate_s) {

cur = ggml_mul(ctx0, cur, gate_s);

cb(cur, "ffn_gate_s", il);

}

} else {

cur = tmp;

}

switch (type_op) {

case LLM_FFN_SILU:

if (gate && type_gate == LLM_FFN_PAR) {

cur = ggml_swiglu_split(ctx0, cur, tmp);

cb(cur, "ffn_swiglu", il);

type_gate = LLM_FFN_SEQ;

} else {

cur = ggml_silu(ctx0, cur);

cb(cur, "ffn_silu", il);

} break;

case LLM_FFN_GELU:

if (gate && type_gate == LLM_FFN_PAR) {

cur = ggml_geglu_split(ctx0, cur, tmp);

cb(cur, "ffn_geglu", il);

type_gate = LLM_FFN_SEQ;

} else {

cur = ggml_gelu(ctx0, cur);

cb(cur, "ffn_gelu", il);

if (act_scales != NULL) {

cur = ggml_div(ctx0, cur, act_scales);

cb(cur, "ffn_act", il);

}

} break;

case LLM_FFN_RELU:

if (gate && type_gate == LLM_FFN_PAR) {

cur = ggml_reglu_split(ctx0, cur, tmp);

cb(cur, "ffn_reglu", il);

type_gate = LLM_FFN_SEQ;

} else {

cur = ggml_relu(ctx0, cur);

cb(cur, "ffn_relu", il);

} break;

case LLM_FFN_RELU_SQR:

{

cur = ggml_relu(ctx0, cur);

cb(cur, "ffn_relu", il);

cur = ggml_sqr(ctx0, cur);

cb(cur, "ffn_sqr(relu)", il);

} break;

case LLM_FFN_SWIGLU:

{

cur = ggml_swiglu(ctx0, cur);

cb(cur, "ffn_swiglu", il);

} break;

case LLM_FFN_GEGLU:

{

cur = ggml_geglu(ctx0, cur);

cb(cur, "ffn_geglu", il);

} break;

case LLM_FFN_REGLU:

{

cur = ggml_reglu(ctx0, cur);

cb(cur, "ffn_reglu", il);

} break;

default:

GGML_ABORT("fatal error");

}

if (gate && type_gate == LLM_FFN_PAR) {

cur = ggml_mul(ctx0, cur, tmp);

cb(cur, "ffn_gate_par", il);

}

if (down) {

cur = build_lora_mm(down, cur);

if (arch == LLM_ARCH_GLM4 || arch == LLM_ARCH_GLM4_MOE) {

// GLM4 and GLM4_MOE seem to have numerical issues with half-precision accumulators

ggml_mul_mat_set_prec(cur, GGML_PREC_F32);

}

}

if (down_b) {

cb(cur, "ffn_down", il);

}

if (down_b) {

cur = ggml_add(ctx0, cur, down_b);

}

if (down_s) {

cur = ggml_mul(ctx0, cur, down_s);

cb(cur, "ffn_down_s", il);

}

return cur;

}

ggml_tensor * llm_graph_context::build_moe_ffn(

ggml_tensor * cur,

ggml_tensor * gate_inp,

ggml_tensor * up_exps,

ggml_tensor * gate_exps,

ggml_tensor * down_exps,

ggml_tensor * exp_probs_b,

int64_t n_expert,

int64_t n_expert_used,

llm_ffn_op_type type_op,

bool norm_w,

bool scale_w,

float w_scale,

llama_expert_gating_func_type gating_op,

int il,

ggml_tensor * probs_in) const {

return build_moe_ffn(

cur,

gate_inp, /* gate_inp_b */ nullptr,

up_exps, /* up_exps_b */ nullptr,

gate_exps, /* gate_exps_b */ nullptr,

down_exps, /* down_exps_b */ nullptr,

exp_probs_b,

n_expert,

n_expert_used,

type_op,

norm_w,

scale_w,

w_scale,

gating_op,

il,

probs_in

);

}

ggml_tensor * llm_graph_context::build_moe_ffn(

ggml_tensor * cur,