#include "llama-impl.h"

#include "llama-chat.h"

#include "llama-mmap.h"

#include "llama-context.h"

#include "llama-vocab.h"

#include "llama-sampling.h"

#include "llama-kv-cache.h"

#include "llama-model-loader.h"

#include "llama-model.h"

#include "ggml.h"

#include "ggml-alloc.h"

#include "ggml-backend.h"

#include "ggml-cpp.h"

#include <algorithm>

#include <array>

#include <cassert>

#include <cfloat>

#include <cmath>

#include <cstddef>

#include <cstdint>

#include <cstdio>

#include <cstring>

#include <ctime>

#include <functional>

#if defined(_MSC_VER)

#pragma warning(disable: 4244 4267) // possible loss of data

#endif

// Returns 0 on success, -1 on error, and -2 on cancellation via llama_progress_callback

static int llama_model_load(const std::string & fname, std::vector<std::string> & splits, llama_model & model, llama_model_params & params) {

// loading time will be recalculated after the first eval, so

// we take page faults deferred by mmap() into consideration

model.t_load_us = 0;

time_meas tm(model.t_load_us);

model.t_start_us = tm.t_start_us;

try {

llama_model_loader ml(fname, splits, params.use_mmap, params.check_tensors, params.kv_overrides);

ml.print_info();

model.hparams.vocab_only = params.vocab_only;

try {

model.load_arch(ml);

} catch(const std::exception & e) {

throw std::runtime_error("error loading model architecture: " + std::string(e.what()));

}

try {

model.load_hparams(ml);

} catch(const std::exception & e) {

throw std::runtime_error("error loading model hyperparameters: " + std::string(e.what()));

}

try {

model.load_vocab(ml);

} catch(const std::exception & e) {

throw std::runtime_error("error loading model vocabulary: " + std::string(e.what()));

}

model.load_stats(ml);

model.print_info();

if (params.vocab_only) {

LLAMA_LOG_INFO("%s: vocab only - skipping tensors\n", __func__);

return 0;

}

if (!model.load_tensors(ml)) {

return -2;

}

} catch (const std::exception & err) {

LLAMA_LOG_ERROR("%s: error loading model: %s\n", __func__, err.what());

return -1;

}

return 0;

}

//

// llm_build

//

using llm_build_cb = std::function<void(struct ggml_tensor * cur, const char * name, int nl)>;

enum llm_ffn_op_type {

LLM_FFN_SILU,

LLM_FFN_GELU,

LLM_FFN_RELU,

LLM_FFN_RELU_SQR,

LLM_FFN_SWIGLU,

};

enum llm_ffn_gate_type {

LLM_FFN_SEQ,

LLM_FFN_PAR, // ffn_gate is parallel to ffn_up

};

enum llm_norm_type {

LLM_NORM,

LLM_NORM_RMS,

LLM_NORM_GROUP,

};

static struct ggml_tensor * llm_build_inp_embd(

struct ggml_context * ctx,

struct llama_context & lctx,

const llama_hparams & hparams,

const llama_ubatch & ubatch,

struct ggml_tensor * tok_embd,

const llm_build_cb & cb) {

const int64_t n_embd = hparams.n_embd;

struct ggml_tensor * inpL;

if (ubatch.token) {

lctx.inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ubatch.n_tokens);

cb(lctx.inp_tokens, "inp_tokens", -1);

ggml_set_input(lctx.inp_tokens);

inpL = ggml_get_rows(ctx, tok_embd, lctx.inp_tokens);

// apply lora for embedding tokens if needed

for (auto & it : lctx.lora) {

struct llama_adapter_lora_weight * lw = it.first->get_weight(tok_embd);

if (lw == nullptr) {

continue;

}

const float adapter_scale = it.second;

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

struct ggml_tensor * inpL_delta = ggml_scale(ctx, ggml_mul_mat(

ctx, lw->b, // non-transposed lora_b

ggml_get_rows(ctx, lw->a, lctx.inp_tokens)

), scale);

inpL = ggml_add(ctx, inpL, inpL_delta);

}

} else {

lctx.inp_embd = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, ubatch.n_tokens);

inpL = lctx.inp_embd;

ggml_set_input(lctx.inp_embd);

}

// For Granite architecture

if (hparams.f_embedding_scale != 0.0f) {

inpL = ggml_scale(ctx, inpL, hparams.f_embedding_scale);

}

cb(inpL, "inp_embd", -1);

return inpL;

}

static void llm_build_kv_store(

struct ggml_context * ctx,

const llama_hparams & hparams,

const llama_cparams & cparams,

const llama_kv_cache & kv,

struct ggml_cgraph * graph,

struct ggml_tensor * k_cur,

struct ggml_tensor * v_cur,

int32_t n_tokens,

int32_t kv_head,

const llm_build_cb & cb,

int64_t il) {

const int64_t n_ctx = cparams.n_ctx;

const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);

const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);

GGML_ASSERT(kv.size == n_ctx);

struct ggml_tensor * k_cache_view = ggml_view_1d(ctx, kv.k_l[il], n_tokens*n_embd_k_gqa, ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa)*kv_head);

cb(k_cache_view, "k_cache_view", il);

// note: storing RoPE-ed version of K in the KV cache

ggml_build_forward_expand(graph, ggml_cpy(ctx, k_cur, k_cache_view));

assert(v_cur->ne[0] == n_embd_v_gqa && v_cur->ne[1] == n_tokens);

struct ggml_tensor * v_cache_view = nullptr;

if (cparams.flash_attn) {

v_cache_view = ggml_view_1d(ctx, kv.v_l[il], n_tokens*n_embd_v_gqa, ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa)*kv_head);

} else {

// note: the V cache is transposed when not using flash attention

v_cache_view = ggml_view_2d(ctx, kv.v_l[il], n_tokens, n_embd_v_gqa,

( n_ctx)*ggml_element_size(kv.v_l[il]),

(kv_head)*ggml_element_size(kv.v_l[il]));

v_cur = ggml_transpose(ctx, v_cur);

}

cb(v_cache_view, "v_cache_view", il);

ggml_build_forward_expand(graph, ggml_cpy(ctx, v_cur, v_cache_view));

}

// do mat_mul, while optionally apply lora

static struct ggml_tensor * llm_build_lora_mm(

struct llama_context & lctx,

struct ggml_context * ctx0,

struct ggml_tensor * w,

struct ggml_tensor * cur) {

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

for (auto & it : lctx.lora) {

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

if (lw == nullptr) {

continue;

}

const float adapter_scale = it.second;

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

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

}

// do mat_mul_id, while optionally apply lora

static struct ggml_tensor * llm_build_lora_mm_id(

struct llama_context & lctx,

struct ggml_context * ctx0,

struct ggml_tensor * w, // struct ggml_tensor * as

struct ggml_tensor * cur, // struct ggml_tensor * b

struct ggml_tensor * ids) {

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

for (auto & it : lctx.lora) {

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

if (lw == nullptr) {

continue;

}

const float alpha = it.first->alpha;

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

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

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

}

static struct ggml_tensor * llm_build_norm(

struct ggml_context * ctx,

struct ggml_tensor * cur,

const llama_hparams & hparams,

struct ggml_tensor * mw,

struct ggml_tensor * mb,

llm_norm_type type,

const llm_build_cb & cb,

int il) {

switch (type) {

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

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

case LLM_NORM_GROUP:

{

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

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

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

} break;

}

if (mw || mb) {

cb(cur, "norm", il);

}

if (mw) {

cur = ggml_mul(ctx, cur, mw);

if (mb) {

cb(cur, "norm_w", il);

}

}

if (mb) {

cur = ggml_add(ctx, cur, mb);

}

return cur;

}

static struct ggml_tensor * llm_build_ffn(

struct ggml_context * ctx,

struct llama_context & lctx,

struct ggml_tensor * cur,

struct ggml_tensor * up,

struct ggml_tensor * up_b,

struct ggml_tensor * up_s,

struct ggml_tensor * gate,

struct ggml_tensor * gate_b,

struct ggml_tensor * gate_s,

struct ggml_tensor * down,

struct ggml_tensor * down_b,

struct ggml_tensor * down_s,

struct ggml_tensor * act_scales,

llm_ffn_op_type type_op,

llm_ffn_gate_type type_gate,

const llm_build_cb & cb,

int il) {

struct ggml_tensor * tmp = up ? llm_build_lora_mm(lctx, ctx, up, cur) : cur;

cb(tmp, "ffn_up", il);

if (up_b) {

tmp = ggml_add(ctx, tmp, up_b);

cb(tmp, "ffn_up_b", il);

}

if (up_s) {

tmp = ggml_mul(ctx, tmp, up_s);

cb(tmp, "ffn_up_s", il);

}

if (gate) {

switch (type_gate) {

case LLM_FFN_SEQ:

{

cur = llm_build_lora_mm(lctx, ctx, gate, tmp);

cb(cur, "ffn_gate", il);

} break;

case LLM_FFN_PAR:

{

cur = llm_build_lora_mm(lctx, ctx, gate, cur);

cb(cur, "ffn_gate", il);

} break;

}

if (gate_b) {

cur = ggml_add(ctx, cur, gate_b);

cb(cur, "ffn_gate_b", il);

}

if (gate_s) {

cur = ggml_mul(ctx, cur, gate_s);

cb(cur, "ffn_gate_s", il);

}

} else {

cur = tmp;

}

switch (type_op) {

case LLM_FFN_SILU:

{

cur = ggml_silu(ctx, cur);

cb(cur, "ffn_silu", il);

} break;

case LLM_FFN_GELU:

{

cur = ggml_gelu(ctx, cur);

cb(cur, "ffn_gelu", il);

if (act_scales != NULL) {

cur = ggml_div(ctx, cur, act_scales);

cb(cur, "ffn_act", il);

}

} break;

case LLM_FFN_RELU:

{

cur = ggml_relu(ctx, cur);

cb(cur, "ffn_relu", il);

} break;

case LLM_FFN_RELU_SQR:

{

cur = ggml_relu(ctx, cur);

cb(cur, "ffn_relu", il);

cur = ggml_sqr(ctx, cur);

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

} break;

case LLM_FFN_SWIGLU:

{

// Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf

int64_t split_point = cur->ne[0] / 2;

struct ggml_tensor * x0 = ggml_cont(ctx, ggml_view_2d(ctx, cur, split_point, cur->ne[1], cur->nb[1], 0));

struct ggml_tensor * x1 = ggml_cont(ctx, ggml_view_2d(ctx, cur, split_point, cur->ne[1], cur->nb[1], split_point * ggml_element_size(cur)));

x0 = ggml_silu(ctx, x0);

cb(cur, "ffn_silu", il);

cur = ggml_mul(ctx, x0, x1);

cb(cur, "ffn_mul", il);

} break;

}

if (type_gate == LLM_FFN_PAR) {

cur = ggml_mul(ctx, cur, tmp);

cb(cur, "ffn_gate_par", il);

}

if (down) {

cur = llm_build_lora_mm(lctx, ctx, down, cur);

}

if (down_b) {

cb(cur, "ffn_down", il);

}

if (down_b) {

cur = ggml_add(ctx, cur, down_b);

}

if (down_s) {

cur = ggml_mul(ctx, cur, down_s);

cb(cur, "ffn_down_s", il);

}

return cur;

}

static struct ggml_tensor * llm_build_moe_ffn(

struct ggml_context * ctx,

struct llama_context & lctx,

struct ggml_tensor * cur,

struct ggml_tensor * gate_inp,

struct ggml_tensor * up_exps,

struct ggml_tensor * gate_exps,

struct ggml_tensor * down_exps,

struct 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,

const llm_build_cb & cb,

int il) {

int64_t n_embd = cur->ne[0];

int64_t n_tokens = cur->ne[1];

ggml_tensor * logits = llm_build_lora_mm(lctx, ctx, gate_inp, cur); // [n_expert, n_tokens]

cb(logits, "ffn_moe_logits", il);

ggml_tensor * probs = nullptr;

switch (gating_op) {

case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX:

{

probs = ggml_soft_max(ctx, logits); // [n_expert, n_tokens]

} break;

case LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID:

{

probs = ggml_sigmoid(ctx, logits); // [n_expert, n_tokens]

} break;

default:

GGML_ABORT("fatal error");

}

cb(probs, "ffn_moe_probs", il);

// add experts selection bias - introduced in DeepSeek V3

// leave probs unbiased as it's later used to get expert weights

ggml_tensor * selection_probs = probs;

if (exp_probs_b != nullptr) {

selection_probs = ggml_add(ctx, probs, exp_probs_b);

cb(selection_probs, "ffn_moe_probs_biased", il);

}

// select experts

ggml_tensor * selected_experts = ggml_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens]

cb(selected_experts->src[0], "ffn_moe_argsort", il);

cb(selected_experts, "ffn_moe_topk", il);

ggml_tensor * weights = ggml_get_rows(ctx,

ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); // [1, n_expert_used, n_tokens]

cb(weights, "ffn_moe_weights", il);

if (norm_w) {

weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens);

ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); // [1, n_tokens]

cb(weights_sum, "ffn_moe_weights_sum", il);

weights = ggml_div(ctx, weights, weights_sum); // [n_expert_used, n_tokens]

cb(weights, "ffn_moe_weights_norm", il);

weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens);

}

if (scale_w) {

weights = ggml_scale(ctx, weights, w_scale);

cb(weights, "ffn_moe_weights_scaled", il);

}

cur = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens);

ggml_tensor * up = llm_build_lora_mm_id(lctx, ctx, up_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]

cb(up, "ffn_moe_up", il);

ggml_tensor * gate = llm_build_lora_mm_id(lctx, ctx, gate_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]

cb(gate, "ffn_moe_gate", il);

switch (type_op) {

case LLM_FFN_SILU:

{

gate = ggml_silu(ctx, gate);

cb(gate, "ffn_moe_silu", il);

} break;

case LLM_FFN_GELU:

{

gate = ggml_gelu(ctx, gate);

cb(gate, "ffn_moe_gelu", il);

} break;

default:

GGML_ABORT("fatal error");

}

ggml_tensor * par = ggml_mul(ctx, up, gate); // [n_ff, n_expert_used, n_tokens]

cb(par, "ffn_moe_gate_par", il);

ggml_tensor * experts = llm_build_lora_mm_id(lctx, ctx, down_exps, par, selected_experts); // [n_embd, n_expert_used, n_tokens]

cb(experts, "ffn_moe_down", il);

experts = ggml_mul(ctx, experts, weights);

// aggregate experts

ggml_tensor * moe_out = nullptr;

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

ggml_tensor * cur_expert = ggml_view_2d(ctx, experts, n_embd, n_tokens,

experts->nb[2], i*experts->nb[1]);

if (i == 0) {

moe_out = cur_expert;

} else {

moe_out = ggml_add(ctx, moe_out, cur_expert);

}

}

if (n_expert_used == 1) {

// avoid returning a non-contiguous tensor

moe_out = ggml_cont(ctx, moe_out);

}

return moe_out;

}

static struct ggml_tensor * llm_build_kqv(

struct ggml_context * ctx,

struct llama_context & lctx,

const llama_kv_cache & kv,

struct ggml_cgraph * graph,

struct ggml_tensor * wo,

struct ggml_tensor * wo_b,

struct ggml_tensor * q_cur,

struct ggml_tensor * kq_mask,

int32_t n_tokens,

int32_t n_kv,

float kq_scale,

const llm_build_cb & cb,

int il) {

const llama_model & model = lctx.model;

const llama_hparams & hparams = lctx.model.hparams;

const llama_cparams & cparams = lctx.cparams;

const int64_t n_ctx = cparams.n_ctx;

const int64_t n_head = hparams.n_head(il);

const int64_t n_head_kv = hparams.n_head_kv(il);

const int64_t n_embd_head_k = hparams.n_embd_head_k;

const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);

const int64_t n_embd_head_v = hparams.n_embd_head_v;

const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);

struct ggml_tensor * q = ggml_permute(ctx, q_cur, 0, 2, 1, 3);

cb(q, "q", il);

struct ggml_tensor * k =

ggml_view_3d(ctx, kv.k_l[il],

n_embd_head_k, n_kv, n_head_kv,

ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa),

ggml_row_size(kv.k_l[il]->type, n_embd_head_k),

0);

cb(k, "k", il);

struct ggml_tensor * cur;

if (cparams.flash_attn) {

GGML_UNUSED(model);

GGML_UNUSED(n_ctx);

// split cached v into n_head heads (not transposed)

struct ggml_tensor * v =

ggml_view_3d(ctx, kv.v_l[il],

n_embd_head_v, n_kv, n_head_kv,

ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa),

ggml_row_size(kv.v_l[il]->type, n_embd_head_v),

0);

cb(v, "v", il);

cur = ggml_flash_attn_ext(ctx, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,

hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);

ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);

cur = ggml_reshape_2d(ctx, cur, n_embd_head_v*n_head, n_tokens);

} else {

struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q);

cb(kq, "kq", il);

// note: this op tends to require high floating point range

// while for some models F16 is enough, for others it is not, so we default to F32 here

ggml_mul_mat_set_prec(kq, GGML_PREC_F32);

if (model.arch == LLM_ARCH_GROK) {

// need to do the following:

// multiply by attn_output_multiplyer of 0.08838834764831845

// and then :

// kq = 30 * tanh(kq / 30)

// before the softmax below

kq = ggml_tanh(ctx, ggml_scale(ctx, kq, 0.08838834764831845f/30.0f));

kq = ggml_scale(ctx, kq, 30);

}

if (hparams.attn_soft_cap) {

kq = ggml_scale(ctx, kq, 1.0f / hparams.f_attn_logit_softcapping);

kq = ggml_tanh(ctx, kq);

kq = ggml_scale(ctx, kq, hparams.f_attn_logit_softcapping);

}

kq = ggml_soft_max_ext(ctx, kq, kq_mask, kq_scale, hparams.f_max_alibi_bias);

cb(kq, "kq_soft_max_ext", il);

GGML_ASSERT(kv.size == n_ctx);

// split cached v into n_head heads

struct ggml_tensor * v =

ggml_view_3d(ctx, kv.v_l[il],

n_kv, n_embd_head_v, n_head_kv,

ggml_element_size(kv.v_l[il])*n_ctx,

ggml_element_size(kv.v_l[il])*n_ctx*n_embd_head_v,

0);

cb(v, "v", il);

struct ggml_tensor * kqv = ggml_mul_mat(ctx, v, kq);

cb(kqv, "kqv", il);

struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 0, 2, 1, 3);

cb(kqv_merged, "kqv_merged", il);

cur = ggml_cont_2d(ctx, kqv_merged, n_embd_head_v*n_head, n_tokens);

cb(cur, "kqv_merged_cont", il);

}

ggml_build_forward_expand(graph, cur);

if (wo) {

cur = llm_build_lora_mm(lctx, ctx, wo, cur);

}

if (wo_b) {

cb(cur, "kqv_wo", il);

}

if (wo_b) {

cur = ggml_add(ctx, cur, wo_b);

}

return cur;

}

static struct ggml_tensor * llm_build_kv(

struct ggml_context * ctx,

struct llama_context & lctx,

const llama_kv_cache & kv,

struct ggml_cgraph * graph,

struct ggml_tensor * wo,

struct ggml_tensor * wo_b,

struct ggml_tensor * k_cur,

struct ggml_tensor * v_cur,

struct ggml_tensor * q_cur,

struct ggml_tensor * kq_mask,

int32_t n_tokens,

int32_t kv_head,

int32_t n_kv,

float kq_scale,

const llm_build_cb & cb,

int il) {

const llama_hparams & hparams = lctx.model.hparams;

const llama_cparams & cparams = lctx.cparams;

// these nodes are added to the graph together so that they are not reordered

// by doing so, the number of splits in the graph is reduced

ggml_build_forward_expand(graph, q_cur);

ggml_build_forward_expand(graph, k_cur);

ggml_build_forward_expand(graph, v_cur);

llm_build_kv_store(ctx, hparams, cparams, kv, graph, k_cur, v_cur, n_tokens, kv_head, cb, il);

struct ggml_tensor * cur;

cur = llm_build_kqv(ctx, lctx, kv, graph, wo, wo_b, q_cur, kq_mask, n_tokens, n_kv, kq_scale, cb, il);

cb(cur, "kqv_out", il);

return cur;

}

static struct ggml_tensor * llm_build_copy_mask_state(

struct ggml_context * ctx,

struct ggml_cgraph * graph,

struct ggml_tensor * s,

struct ggml_tensor * state_copy,

struct ggml_tensor * state_mask,

int32_t n_state,

int32_t kv_size,

int32_t kv_head,

int32_t n_kv,

int32_t n_seqs) {

struct ggml_tensor * states = ggml_reshape_2d(ctx, s, n_state, kv_size);

// copy states

// NOTE: assuming the copy destinations are ALL contained between kv_head and kv_head + n_kv

// this shrinks the tensors's ne[1] to n_kv

states = ggml_get_rows(ctx, states, state_copy);

// clear states of sequences which are starting at the beginning of this batch

// FIXME: zero-out NANs?

states = ggml_mul(ctx, states, state_mask);

// copy states which won't be changed further (between n_seqs and n_kv)

ggml_build_forward_expand(graph,

ggml_cpy(ctx,

ggml_view_1d(ctx, states, n_state*(n_kv - n_seqs), n_seqs*n_state*ggml_element_size(states)),

ggml_view_1d(ctx, s, n_state*(n_kv - n_seqs), (kv_head + n_seqs)*n_state*ggml_element_size(s))));

// the part of the states that will be used and modified

return ggml_view_2d(ctx, states, n_state, n_seqs, states->nb[1], 0);

}

// TODO: split

static struct ggml_tensor * llm_build_mamba(

struct ggml_context * ctx,

struct llama_context & lctx,

const llama_ubatch & ubatch,

struct ggml_cgraph * graph,

struct ggml_tensor * cur,

struct ggml_tensor * state_copy,

struct ggml_tensor * state_mask,

int32_t kv_head,

int32_t n_kv,

const llm_build_cb & cb,

int il) {

const llama_model & model = lctx.model;

const llama_hparams & hparams = model.hparams;

const llama_kv_cache & kv = lctx.kv_self;

const int64_t d_conv = hparams.ssm_d_conv;

const int64_t d_inner = hparams.ssm_d_inner;

const int64_t d_state = hparams.ssm_d_state;

const int64_t dt_rank = hparams.ssm_dt_rank;

const int64_t n_seqs = ubatch.n_seqs;

// Some variants of Mamba arch (e.g. FalconMamba do apply layer norm on B and Dt layers)

const bool ssm_dt_b_c_rms = hparams.ssm_dt_b_c_rms;

// Use the same RMS norm as the final layer norm

const float norm_rms_eps = hparams.f_norm_rms_eps;

const int64_t n_seq_tokens = ubatch.n_seq_tokens;

GGML_ASSERT(n_seqs != 0);

GGML_ASSERT(ubatch.equal_seqs);

GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);

struct ggml_tensor * conv_states_all = kv.k_l[il];

struct ggml_tensor * ssm_states_all = kv.v_l[il];

// (ab)using the KV cache to store the states

struct ggml_tensor * conv = llm_build_copy_mask_state(ctx,

graph, conv_states_all, state_copy, state_mask,

hparams.n_embd_k_s(), kv.size, kv_head, n_kv, n_seqs);

conv = ggml_reshape_3d(ctx, conv, d_conv - 1, d_inner, n_seqs);

struct ggml_tensor * ssm = llm_build_copy_mask_state(ctx,

graph, ssm_states_all, state_copy, state_mask,

hparams.n_embd_v_s(), kv.size, kv_head, n_kv, n_seqs);

ssm = ggml_reshape_3d(ctx, ssm, d_state, d_inner, n_seqs);

// {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}

cur = ggml_reshape_3d(ctx, cur, cur->ne[0], n_seq_tokens, n_seqs);

// {n_embd, 2*d_inner} @ {n_embd, n_seq_tokens, n_seqs} => {2*d_inner, n_seq_tokens, n_seqs}

struct ggml_tensor * xz = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_in, cur);

// split the above in two

// => {d_inner, n_seq_tokens, n_seqs}

struct ggml_tensor * x = ggml_view_3d(ctx, xz, d_inner, xz->ne[1], xz->ne[2], xz->nb[1], xz->nb[2], 0);

struct ggml_tensor * z = ggml_view_3d(ctx, xz, d_inner, xz->ne[1], xz->ne[2], xz->nb[1], xz->nb[2], d_inner*ggml_element_size(xz));

// conv

{

// => {d_conv - 1 + n_seq_tokens, d_inner, n_seqs}

struct ggml_tensor * conv_x = ggml_concat(ctx, conv, ggml_transpose(ctx, x), 0);

// copy last (d_conv - 1) columns back into the state cache

struct ggml_tensor * last_conv = ggml_view_3d(ctx, conv_x, d_conv - 1, d_inner, n_seqs, conv_x->nb[1], conv_x->nb[2], n_seq_tokens*(conv_x->nb[0]));

ggml_build_forward_expand(graph,

ggml_cpy(ctx, last_conv,

ggml_view_1d(ctx, conv_states_all,

(d_conv - 1)*(d_inner)*(n_seqs),

kv_head*(d_conv - 1)*(d_inner)*ggml_element_size(conv_states_all))));

// 1D convolution

// The equivalent is to make a self-overlapping view of conv_x

// over d_conv columns at each stride in the 3rd dimension,

// then element-wise multiply that with the conv1d weight,

// then sum the elements of each row,

// (the last two steps are a dot product over rows (also doable with mul_mat))

// then permute away the ne[0] dimension,

// and then you're left with the resulting x tensor.

// For simultaneous sequences, all sequences need to have the same length.

x = ggml_ssm_conv(ctx, conv_x, model.layers[il].ssm_conv1d);

// bias

x = ggml_add(ctx, x, model.layers[il].ssm_conv1d_b);

x = ggml_silu(ctx, x);

}

// ssm

{

// {d_inner, dt_rank + 2*d_state} @ {d_inner, n_seq_tokens, n_seqs} => {dt_rank + 2*d_state, n_seq_tokens, n_seqs}

struct ggml_tensor * x_db = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_x, x);

// split

struct ggml_tensor * dt = ggml_view_3d(ctx, x_db, dt_rank, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], 0);

struct ggml_tensor * B = ggml_view_3d(ctx, x_db, d_state, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], ggml_element_size(x_db)*dt_rank);

struct ggml_tensor * C = ggml_view_3d(ctx, x_db, d_state, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], ggml_element_size(x_db)*(dt_rank+d_state));

// Some Mamba variants (e.g. FalconMamba) apply RMS norm in B, C & Dt layers

if (ssm_dt_b_c_rms) {

dt = ggml_rms_norm(ctx, dt, norm_rms_eps);

B = ggml_rms_norm(ctx, B, norm_rms_eps);

C = ggml_rms_norm(ctx, C, norm_rms_eps);

}

// {dt_rank, d_inner} @ {dt_rank, n_seq_tokens, n_seqs} => {d_inner, n_seq_tokens, n_seqs}

dt = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_dt, dt);

dt = ggml_add(ctx, dt, model.layers[il].ssm_dt_b);

// Custom operator to optimize the parallel associative scan

// as described in the Annex D of the Mamba paper.

// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}

struct ggml_tensor * y_ssm = ggml_ssm_scan(ctx, ssm, x, dt, model.layers[il].ssm_a, B, C);

// store last states

ggml_build_forward_expand(graph,

ggml_cpy(ctx,

ggml_view_1d(ctx, y_ssm, d_state*d_inner*n_seqs, x->nb[3]),

ggml_view_1d(ctx, ssm_states_all, d_state*d_inner*n_seqs, kv_head*d_state*d_inner*ggml_element_size(ssm_states_all))));

struct ggml_tensor * y = ggml_view_3d(ctx, y_ssm, d_inner, n_seq_tokens, n_seqs, x->nb[1], x->nb[2], 0);

// TODO: skip computing output earlier for unused tokens

// {d_inner, n_seq_tokens, n_seqs} * {d_inner} => {d_inner, n_seq_tokens, n_seqs}

y = ggml_add(ctx, y, ggml_mul(ctx, x, model.layers[il].ssm_d));

y = ggml_mul(ctx, y, ggml_silu(ctx, ggml_cont(ctx, z)));

// {d_inner, n_embd} @ {d_inner, n_seq_tokens, n_seqs} => {n_embd, n_seq_tokens, n_seqs}

cur = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_out, y);

}

// {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}

cur = ggml_reshape_2d(ctx, cur, cur->ne[0], n_seq_tokens * n_seqs);

cb(cur, "mamba_out", il);

return cur;

}

static struct ggml_tensor * llm_build_rwkv6_time_mix(

struct llama_context & lctx,

struct ggml_context * ctx,

const struct llama_layer * layer,

struct ggml_tensor * cur,

struct ggml_tensor * x_prev,

struct ggml_tensor ** wkv_state,

size_t wkv_head_size,

size_t head_count_kv) {

size_t n_embd = cur->ne[0];

size_t n_seq_tokens = cur->ne[1];

size_t n_seqs = cur->ne[2];

size_t head_size = wkv_head_size;

size_t head_count = n_embd / head_size;

size_t n_tokens = n_seqs * n_seq_tokens;

bool is_qrwkv = layer->time_mix_first == nullptr;

struct ggml_tensor * sx = ggml_sub(ctx, x_prev, cur);

sx = ggml_reshape_2d(ctx, sx, n_embd, n_tokens);

cur = ggml_reshape_2d(ctx, cur, n_embd, n_tokens);

struct ggml_tensor * xxx = ggml_add(ctx, ggml_mul(ctx, sx, layer->time_mix_lerp_x), cur);

xxx = ggml_reshape_4d(

ctx,

ggml_tanh(

ctx,

ggml_mul_mat(ctx, layer->time_mix_w1, xxx)

),

layer->time_mix_w1->ne[1] / 5, 1, 5, n_tokens

);

xxx = ggml_cont(ctx, ggml_permute(ctx, xxx, 0, 1, 3, 2));

xxx = ggml_mul_mat(

ctx,

ggml_reshape_4d(

ctx,

layer->time_mix_w2,

layer->time_mix_w2->ne[0], layer->time_mix_w2->ne[1], 1, 5

),

xxx

);

struct ggml_tensor *xw, *xk, *xv, *xr, *xg;

if (layer->time_mix_lerp_fused) {

// fusing these weights makes some performance improvement

sx = ggml_reshape_3d(ctx, sx, n_embd, 1, n_tokens);

cur = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens);

xxx = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xxx, layer->time_mix_lerp_fused), sx), cur);

xw = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], 0);

xk = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));

xv = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));

xr = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));

xg = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));

} else {

// for backward compatibility

xw = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], 0);

xk = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));

xv = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));

xr = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));

xg = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));

xw = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xw, layer->time_mix_lerp_w), sx), cur);

xk = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xk, layer->time_mix_lerp_k), sx), cur);

xv = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xv, layer->time_mix_lerp_v), sx), cur);

xr = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xr, layer->time_mix_lerp_r), sx), cur);

xg = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xg, layer->time_mix_lerp_g), sx), cur);

}

struct ggml_tensor * r = llm_build_lora_mm(lctx, ctx, layer->time_mix_receptance, xr);

struct ggml_tensor * k = llm_build_lora_mm(lctx, ctx, layer->time_mix_key, xk);

struct ggml_tensor * v = llm_build_lora_mm(lctx, ctx, layer->time_mix_value, xv);

if (layer->time_mix_receptance_b) {

r = ggml_add(ctx, r, layer->time_mix_receptance_b);

}

if (layer->time_mix_key_b) {

k = ggml_add(ctx, k, layer->time_mix_key_b);

}

if (layer->time_mix_value_b) {

v = ggml_add(ctx, v, layer->time_mix_value_b);

}

struct ggml_tensor * g = llm_build_lora_mm(lctx, ctx, layer->time_mix_gate, xg);

if (is_qrwkv) {

g = ggml_sigmoid(ctx, g);

} else {

g = ggml_silu(ctx, g);

}

if (head_count_kv != head_count) {

GGML_ASSERT(head_count % head_count_kv == 0);

k = ggml_reshape_4d(ctx, k, head_size, 1, head_count_kv, n_tokens);

v = ggml_reshape_4d(ctx, v, head_size, 1, head_count_kv, n_tokens);

struct ggml_tensor * tmp = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_size, head_count / head_count_kv, head_count_kv, n_tokens);

k = ggml_repeat(ctx, k, tmp);

v = ggml_repeat(ctx, v, tmp);

}

k = ggml_reshape_3d(ctx, k, head_size, head_count, n_tokens);

v = ggml_reshape_3d(ctx, v, head_size, head_count, n_tokens);

r = ggml_reshape_3d(ctx, r, head_size, head_count, n_tokens);

struct ggml_tensor * w = ggml_mul_mat(

ctx,

layer->time_mix_decay_w2,

ggml_tanh(

ctx,

ggml_mul_mat(ctx, layer->time_mix_decay_w1, xw)

)

);

w = ggml_add(ctx, w, layer->time_mix_decay);

w = ggml_exp(ctx, ggml_neg(ctx, ggml_exp(ctx, w)));

w = ggml_reshape_3d(ctx, w, head_size, head_count, n_tokens);

if (is_qrwkv) {

// k = k * (1 - w)

k = ggml_sub(ctx, k, ggml_mul(ctx, k, w));

}

struct ggml_tensor * wkv_output;

if (!layer->time_mix_first) {

wkv_output = ggml_gated_linear_attn(ctx, k, v, r, w, *wkv_state, pow(head_size, -0.5f));

} else {

wkv_output = ggml_rwkv_wkv6(ctx, k, v, r, layer->time_mix_first, w, *wkv_state);

}

cur = ggml_view_1d(ctx, wkv_output, n_embd * n_tokens, 0);

*wkv_state = ggml_view_1d(ctx, wkv_output, n_embd * head_size * n_seqs, n_embd * n_tokens * sizeof(float));