CS348Project/llama.cpp
5
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-11-08 10:07:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

enc-dec : compose wip

Browse Source 
ggml-ci
This commit is contained in:
Georgi Gerganov
2025-02-24 15:16:45 +02:00
parent 9cd78f11a1
commit be58e30017
5 changed files with 1002 additions and 404 deletions
Download Patch File Download Diff File Expand all files Collapse all files

777
src/llama-context.cpp
View File

File diff suppressed because it is too large Load Diff

261
src/llama-context.h
View File

@@ -30,8 +30,7 @@ public:
virtual void synchronize() = 0;
virtual const llama_model & get_model() const = 0;
virtual const llama_cparams & get_cparams() const = 0;
virtual const llama_model & get_model() const = 0;
virtual uint32_t n_ctx() const = 0;
virtual uint32_t n_ctx_per_seq() const = 0;
@@ -42,8 +41,6 @@ public:
virtual uint32_t n_threads() const = 0;
virtual uint32_t n_threads_batch() const = 0;
virtual int32_t max_nodes() const = 0;
// self-attention:
// if the context does not have a KV cache, return nullptr
@@ -62,8 +59,6 @@ public:
virtual float * get_embeddings_ith(int32_t i) = 0;
virtual float * get_embeddings_seq(llama_seq_id seq_id) = 0;
virtual int64_t n_pos_per_token() const = 0; // vision
virtual void attach_threadpool(
ggml_threadpool_t threadpool,
ggml_threadpool_t threadpool_batch) = 0;
@@ -190,8 +185,7 @@ protected:
virtual void reserve();
public:
const llama_model & get_model() const override;
const llama_cparams & get_cparams() const override;
const llama_model & get_model() const override;
uint32_t n_ctx() const override;
uint32_t n_ctx_per_seq() const override;
@@ -202,15 +196,9 @@ public:
uint32_t n_threads() const override;
uint32_t n_threads_batch() const override;
int32_t max_nodes() const override;
// self-attention:
// if the context does not have a KV cache, return nullptr
llama_kv_cache * get_kv_self() override;
const llama_kv_cache * get_kv_self() const override;
// if the context does not have a KV cache, noop
void kv_self_update() override;
enum llama_pooling_type pooling_type() const override;
@@ -222,8 +210,6 @@ public:
float * get_embeddings_ith(int32_t i) override;
float * get_embeddings_seq(llama_seq_id seq_id) override;
int64_t n_pos_per_token() const override; // vision
void attach_threadpool(
ggml_threadpool_t threadpool,
ggml_threadpool_t threadpool_batch) override;
@@ -261,6 +247,8 @@ protected:
// input
//
virtual int64_t n_pos_per_token() const; // vision
// when the compute graph is built, it creates the input tensors that it needs
// the contents of the input tensors are set by the input_set() function
@@ -299,6 +287,8 @@ protected:
// graph
//
virtual int32_t graph_max_nodes() const;
// zero-out inputs and create the ctx_compute for the compute graph
virtual ggml_cgraph * graph_init();
@@ -477,11 +467,11 @@ public:
size_t n_token_count) override;
protected:
virtual size_t state_get_data(llama_io_write_i & io);
virtual size_t state_set_data(llama_io_read_i & io);
virtual size_t state_write_data(llama_io_write_i & io);
virtual size_t state_read_data (llama_io_read_i & io);
virtual size_t state_seq_get_data(llama_io_write_i & io, llama_seq_id seq_id);
virtual size_t state_seq_set_data(llama_io_read_i & io, llama_seq_id seq_id);
virtual size_t state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id);
virtual size_t state_seq_read_data (llama_io_read_i & io, llama_seq_id seq_id);
//
// members
@@ -625,39 +615,15 @@ protected:
ggml_context * ctx0,
ggml_cgraph * gf) override;
// =======================================================
// === encoder-decoder ===
//
// TODO: this is temporary here, it will be moved
//
// whether we are computing encoder output or decoder output
bool is_encoding = false;
// output of the encoder part of the encoder-decoder models
std::vector<float> embd_enc;
std::vector<std::set<llama_seq_id>> seq_ids_enc;
struct ggml_tensor * inp_embd_enc; // F32 [n_embd, n_outputs_enc]
struct ggml_tensor * inp_kq_mask_cross; // F32 [n_outputs_enc, n_batch]
ggml_tensor * build_inp_embd_enc(
ggml_context * ctx0) override;
ggml_tensor * build_inp_kq_mask_cross(
ggml_context * ctx0,
int32_t n_tokens) override;
// ======================================================
//
// state save/load
//
size_t state_get_data(llama_io_write_i & io) override;
size_t state_set_data(llama_io_read_i & io) override;
size_t state_write_data(llama_io_write_i & io) override;
size_t state_read_data (llama_io_read_i & io) override;
size_t state_seq_get_data(llama_io_write_i & io, llama_seq_id seq_id) override;
size_t state_seq_set_data(llama_io_read_i & io, llama_seq_id seq_id) override;
size_t state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id) override;
size_t state_seq_read_data (llama_io_read_i & io, llama_seq_id seq_id) override;
private:
//
@@ -767,11 +733,11 @@ protected:
// state save/load
//
size_t state_get_data(llama_io_write_i & io) override;
size_t state_set_data(llama_io_read_i & io) override;
size_t state_write_data(llama_io_write_i & io) override;
size_t state_read_data (llama_io_read_i & io) override;
size_t state_seq_get_data(llama_io_write_i & io, llama_seq_id seq_id) override;
size_t state_seq_set_data(llama_io_read_i & io, llama_seq_id seq_id) override;
size_t state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id) override;
size_t state_seq_read_data (llama_io_read_i & io, llama_seq_id seq_id) override;
private:
//
@@ -782,21 +748,206 @@ private:
llama_kv_cache_recurrent kv_self;
};
// TODO: tmp - need something better
struct llama_cross {
int32_t n_outputs;
float * embd_enc;
std::vector<std::set<llama_seq_id>> seq_ids_enc;
};
class llama_context_enc : public llama_context_base {
public:
using llama_context_base::llama_context_base;
int encode(llama_batch & inp_batch) override;
llama_cross * cross = nullptr;
};
class llama_context_enc_dec : public llama_context_enc {
class llama_context_dec : public llama_context_kv_self {
public:
using llama_context_kv_self::llama_context_kv_self;
protected:
void reserve() override;
//
// input
//
void input_set(const llama_ubatch & ubatch) override;
private:
struct {
ggml_tensor * cross_embd; // F32 [n_embd, n_outputs_enc]
ggml_tensor * cross_kq_mask; // F32 [n_outputs_enc, n_batch]
ggml_tensor * cross_kq_mask_cnv; // F32 [n_outputs_enc, n_batch]
} inp;
protected:
//
// graph
//
ggml_cgraph * graph_init() override;
ggml_tensor * build_inp_cross_embd(
ggml_context * ctx0) override;
void build_attn_inp(
ggml_context * ctx0,
int32_t n_tokens,
bool causal,
bool swa) override;
ggml_tensor * build_attn_cross(
ggml_context * ctx0,
ggml_cgraph * gf,
ggml_tensor * q_cur,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
ggml_tensor * kq_b,
float kq_scale,
int il) override;
public:
llama_cross * cross = nullptr;
};
class llama_context_enc_dec : public llama_context_i {
public:
llama_context_enc_dec(
const llama_model & model,
llama_context_params params);
virtual ~llama_context_enc_dec();
~llama_context_enc_dec();
void init() override;
void synchronize() override;
const llama_model & get_model() const override;
// TODO: the default implementation of these getters calls the corresponding getter of the enc or dec context
// in the future, the public API in llama.h should allow to get references to the context that the user wants
// this will allow to specify the desired context explicitly
// for example:
//
// // this can be an enc-dec context
// llama_context_t ctx = llama_init_from_model(...);
//
// ...
//
// llama_context_t ctx_enc = llama_get_ctx_enc(ctx);
// llama_set_embeddings(ctx_enc, true);
//
// llama_context_t ctx_dec = llama_get_ctx_dec(ctx);
// llama_set_causal_attn(ctx_dec, true);
//
uint32_t n_ctx() const override;
uint32_t n_ctx_per_seq() const override;
uint32_t n_batch() const override;
uint32_t n_ubatch() const override;
uint32_t n_seq_max() const override;
uint32_t n_threads() const override;
uint32_t n_threads_batch() const override;
llama_kv_cache * get_kv_self() override;
const llama_kv_cache * get_kv_self() const override;
void kv_self_update() override;
enum llama_pooling_type pooling_type() const override;
float * get_logits() override;
float * get_logits_ith(int32_t i) override;
float * get_embeddings() override;
float * get_embeddings_ith(int32_t i) override;
float * get_embeddings_seq(llama_seq_id seq_id) override;
void attach_threadpool(
ggml_threadpool_t threadpool,
ggml_threadpool_t threadpool_batch) override;
void detach_threadpool() override;
void set_n_threads(int32_t n_threads, int32_t n_threads_batch) override;
void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data) override;
void set_embeddings (bool value) override;
void set_causal_attn(bool value) override;
void set_adapter_lora(
llama_adapter_lora * adapter,
float scale) override;
bool rm_adapter_lora(
llama_adapter_lora * adapter) override;
void clear_adapter_lora() override;
bool apply_adapter_cvec(
const float * data,
size_t len,
int32_t n_embd,
int32_t il_start,
int32_t il_end) override;
int encode(llama_batch & inp_batch) override;
int decode(llama_batch & inp_batch) override;
//
// perf
//
llama_perf_context_data perf_get_data() const override;
void perf_reset() override;
//
// state save/load
//
size_t state_get_size() override;
size_t state_get_data( uint8_t * dst, size_t size) override;
size_t state_set_data(const uint8_t * src, size_t size) override;
size_t state_seq_get_size(llama_seq_id seq_id) override;
size_t state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size) override;
size_t state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size) override;
bool state_load_file(
const char * filepath,
llama_token * tokens_out,
size_t n_token_capacity,
size_t * n_token_count_out) override;
bool state_save_file(
const char * filepath,
const llama_token * tokens,
size_t n_token_count) override;
size_t state_seq_load_file(
llama_seq_id seq_id,
const char * filepath,
llama_token * tokens_out,
size_t n_token_capacity,
size_t * n_token_count_out) override;
size_t state_seq_save_file(
llama_seq_id seq_id,
const char * filepath,
const llama_token * tokens,
size_t n_token_count) override;
private:
llama_context_kv_self ctx_dec;
std::unique_ptr<llama_context_enc> ctx_enc;
std::unique_ptr<llama_context_dec> ctx_dec;
llama_cross cross;
};
// For internal test use

26
src/llama-graph.cpp
View File

@@ -26,7 +26,29 @@ ggml_tensor * llama_graph_i::build_attn(
return nullptr;
}
ggml_tensor * llama_graph_i::build_inp_embd_enc(
ggml_tensor * llama_graph_i::build_attn_cross(
ggml_context * ctx0,
ggml_cgraph * gf,
ggml_tensor * q_cur,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
ggml_tensor * kq_b,
float kq_scale,
int il) {
GGML_UNUSED(ctx0);
GGML_UNUSED(gf);
GGML_UNUSED(q_cur);
GGML_UNUSED(k_cur);
GGML_UNUSED(v_cur);
GGML_UNUSED(kq_b);
GGML_UNUSED(kq_scale);
GGML_UNUSED(il);
LLAMA_LOG_ERROR("%s: not implemented\n", __func__);
return nullptr;
}
ggml_tensor * llama_graph_i::build_inp_cross_embd(
ggml_context * ctx0) {
GGML_UNUSED(ctx0);
@@ -34,7 +56,7 @@ ggml_tensor * llama_graph_i::build_inp_embd_enc(
return nullptr;
}
ggml_tensor * llama_graph_i::build_inp_kq_mask_cross(
ggml_tensor * llama_graph_i::build_inp_cross_kq_mask(
ggml_context * ctx0,
int32_t n_tokens) {
GGML_UNUSED(ctx0);

14
src/llama-graph.h
View File

@@ -114,10 +114,20 @@ public:
float kq_scale,
int il);
virtual ggml_tensor * build_inp_embd_enc(
virtual ggml_tensor * build_attn_cross(
ggml_context * ctx0,
ggml_cgraph * gf,
ggml_tensor * q_cur,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
ggml_tensor * kq_b,
float kq_scale,
int il);
virtual ggml_tensor * build_inp_cross_embd(
ggml_context * ctx0);
virtual ggml_tensor * build_inp_kq_mask_cross(
virtual ggml_tensor * build_inp_cross_kq_mask(
ggml_context * ctx0,
int32_t n_tokens);

328
src/llama-model.cpp
View File

@@ -3964,16 +3964,16 @@ struct llm_build_context {
}
// TODO: tmp
struct ggml_tensor * build_inp_embd_enc() {
ggml_tensor * cur = lgf->build_inp_embd_enc(ctx0);
struct ggml_tensor * build_inp_cross_embd() {
ggml_tensor * cur = lgf->build_inp_cross_embd(ctx0);
cb(cur, "embd_enc", -1);
return cur;
}
// TODO: tmp
struct ggml_tensor * build_inp_kq_mask_cross() {
ggml_tensor * cur = lgf->build_inp_kq_mask_cross(ctx0, n_tokens);
struct ggml_tensor * build_inp_cross_kq_mask() {
ggml_tensor * cur = lgf->build_inp_cross_kq_mask(ctx0, n_tokens);
cb(cur, "KQ_mask_cross", -1);
return cur;
@@ -4294,6 +4294,42 @@ struct llm_build_context {
return cur;
}
struct ggml_tensor * build_attn_cross(
struct ggml_cgraph * gf,
struct ggml_tensor * wo,
struct ggml_tensor * wo_b,
struct ggml_tensor * q_cur,
struct ggml_tensor * k_cur,
struct ggml_tensor * v_cur,
int32_t n_tokens, // TODO: remove
float kq_scale,
int il) {
GGML_UNUSED(n_tokens);
// 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(gf, q_cur);
ggml_build_forward_expand(gf, k_cur);
ggml_build_forward_expand(gf, v_cur);
ggml_tensor * cur = lgf->build_attn_cross(ctx0, gf, q_cur, k_cur, v_cur, nullptr, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
cur = lgf->build_lora_mm(ctx0, wo, cur);
}
if (wo_b) {
//cb(cur, "kqv_wo", il);
}
if (wo_b) {
cur = ggml_add(ctx0, cur, wo_b);
}
return cur;
}
struct ggml_tensor * build_attn_with_kq_b(
struct ggml_cgraph * gf,
struct ggml_tensor * wo,
@@ -9762,209 +9798,173 @@ struct llm_build_context {
ggml_build_forward_expand(gf, cur);
}
//void build_t5_dec(ggml_cgraph * gf) {
// const int64_t n_embd_head = hparams.n_embd_head_v;
// const int64_t n_embd_gqa = hparams.n_embd_v_gqa();
void build_t5_dec(ggml_cgraph * gf) {
const int64_t n_embd_head = hparams.n_embd_head_v;
//const int64_t n_embd_gqa = hparams.n_embd_v_gqa();
// GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
// struct ggml_tensor * cur;
// struct ggml_tensor * inpL;
struct ggml_tensor * cur;
struct ggml_tensor * inpL;
// inpL = build_inp_embd(model.tok_embd);
inpL = build_inp_embd(model.tok_embd);
// GGML_ASSERT(!lctx.is_encoding);
// GGML_ASSERT(n_outputs_enc > 0 && "call llama_encode() first");
struct ggml_tensor * embd_enc = build_inp_cross_embd();
struct ggml_tensor * pos_bucket_dec = build_pos_bucket();
// struct ggml_tensor * embd_enc = build_inp_embd_enc();
// struct ggml_tensor * pos_bucket_dec = build_pos_bucket(true);
const int64_t n_outputs_enc = embd_enc->ne[1];
// struct ggml_tensor * KQ_mask_dec = build_inp_kq_mask();
// struct ggml_tensor * KQ_mask_cross = build_inp_kq_mask_cross();
lgf->build_attn_inp(ctx0, n_tokens, true, false);
// for (int il = 0; il < n_layer; ++il) {
// struct ggml_tensor * inpSA = inpL;
for (int il = 0; il < n_layer; ++il) {
struct ggml_tensor * inpSA = inpL;
// // norm
// cur = build_norm(inpL,
// model.layers[il].attn_norm, NULL,
// LLM_NORM_RMS, il);
// cb(cur, "attn_norm", il);
// norm
cur = build_norm(inpL,
model.layers[il].attn_norm, NULL,
LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
// // self-attention
// {
// struct ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
// cb(Qcur, "Qcur", il);
// self-attention
{
struct ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
cb(Qcur, "Qcur", il);
// struct ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
// cb(Kcur, "Kcur", il);
struct ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
cb(Kcur, "Kcur", il);
// struct ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
// cb(Vcur, "Vcur", il);
struct ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
cb(Vcur, "Vcur", il);
// build_kv_store(gf, Kcur, Vcur, il);
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
// struct ggml_tensor * k =
// ggml_view_3d(ctx0, kv_self.k_l[il],
// n_embd_head_k, n_kv, n_head_kv,
// ggml_row_size(kv_self.k_l[il]->type, n_embd_k_gqa),
// ggml_row_size(kv_self.k_l[il]->type, n_embd_head_k),
// 0);
// cb(k, "k", il);
struct ggml_tensor * attn_rel_b = model.layers[il].attn_rel_b ? model.layers[il].attn_rel_b : model.layers[0].attn_rel_b;
struct ggml_tensor * kq_b = build_pos_bias(pos_bucket_dec, attn_rel_b);
// struct ggml_tensor * v =
// ggml_view_3d(ctx0, kv_self.v_l[il],
// n_kv, n_embd_head_v, n_head_kv,
// ggml_element_size(kv_self.v_l[il])*n_ctx,
// ggml_element_size(kv_self.v_l[il])*n_ctx*n_embd_head_v,
// 0);
// cb(v, "v", il);
cur = build_attn_with_kq_b(gf,
model.layers[il].wo, model.layers[il].bo,
Qcur, Kcur, Vcur, kq_b, n_tokens, 1.0f, il);
cb(cur, "kqv_out", il);
}
// Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
cur = ggml_add(ctx0, cur, inpSA);
cb(cur, "cross_inp", il);
// struct ggml_tensor * q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
struct ggml_tensor * inpCA = cur;
// struct ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
// cb(kq, "kq", il);
// norm
cur = build_norm(cur,
model.layers[il].attn_norm_cross, NULL,
LLM_NORM_RMS, il);
cb(cur, "attn_norm_cross", il);
// struct ggml_tensor * attn_rel_b = model.layers[il].attn_rel_b ? model.layers[il].attn_rel_b : model.layers[0].attn_rel_b;
// struct ggml_tensor * pos_bias = build_pos_bias(pos_bucket_dec, attn_rel_b);
// struct ggml_tensor * kq_b = ggml_add(ctx0, kq, pos_bias);
// cb(kq_b, "kq_b", il);
// cross-attention
{
struct ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq_cross, cur);
cb(Qcur, "Qcur", il);
// kq = ggml_soft_max_ext(ctx0, kq_b, KQ_mask_dec, 1.0f, hparams.f_max_alibi_bias);
// cb(kq, "kq_soft_max_ext", il);
struct ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk_cross, embd_enc);
cb(Kcur, "Kcur", il);
// struct ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq);
// cb(kqv, "kqv", il);
struct ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv_cross, embd_enc);
cb(Vcur, "Vcur", il);
// struct ggml_tensor * kqv_merged = ggml_permute(ctx0, kqv, 0, 2, 1, 3);
// cb(kqv_merged, "kqv_merged", il);
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_outputs_enc);
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_outputs_enc);
// cur = ggml_cont_2d(ctx0, kqv_merged, n_embd_gqa, n_tokens);
// cb(cur, "kqv_merged_cont", il);
cur = build_attn_cross(gf,
model.layers[il].wo_cross, nullptr,
Qcur, Kcur, Vcur, n_tokens, 1.0f, il);
cb(cur, "kqv_out", il);
// ggml_build_forward_expand(gf, cur);
//struct ggml_tensor * q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
//struct ggml_tensor * k = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 0, 2, 1, 3));
// cur = build_lora_mm(model.layers[il].wo, cur);
// cb(cur, "kqv_out", il);
// }
//struct ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
//cb(kq, "kq", il);
// cur = ggml_add(ctx0, cur, inpSA);
// cb(cur, "cross_inp", il);
//kq = ggml_soft_max_ext(ctx0, kq, KQ_mask_cross, 1.0f, hparams.f_max_alibi_bias);
//cb(kq, "kq_soft_max_ext", il);
// struct ggml_tensor * inpCA = cur;
//struct ggml_tensor * v = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_reshape_2d(ctx0, Vcur, n_embd_gqa, n_outputs_enc)));
//cb(v, "v", il);
// // norm
// cur = build_norm(cur,
// model.layers[il].attn_norm_cross, NULL,
// LLM_NORM_RMS, il);
// cb(cur, "attn_norm_cross", il);
//struct ggml_tensor * kqv = ggml_mul_mat(ctx0, ggml_reshape_3d(ctx0, v, n_outputs_enc, n_embd_head, n_head_kv), kq);
//cb(kqv, "kqv", il);
// // cross-attention
// {
// struct ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq_cross, cur);
// cb(Qcur, "Qcur", il);
//struct ggml_tensor * kqv_merged = ggml_permute(ctx0, kqv, 0, 2, 1, 3);
//cb(kqv_merged, "kqv_merged", il);
// struct ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk_cross, embd_enc);
// cb(Kcur, "Kcur", il);
//cur = ggml_cont_2d(ctx0, kqv_merged, n_embd_gqa, n_tokens);
//cb(cur, "kqv_merged_cont", il);
// struct ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv_cross, embd_enc);
// cb(Vcur, "Vcur", il);
//ggml_build_forward_expand(gf, cur);
// Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
// Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_outputs_enc);
//cur = build_lora_mm(model.layers[il].wo_cross, cur);
//cb(cur, "kqv_out", il);
}
// struct ggml_tensor * q = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
// struct ggml_tensor * k = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 0, 2, 1, 3));
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
inpCA = ggml_get_rows(ctx0, inpCA, inp_out_ids);
}
// struct ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
// cb(kq, "kq", il);
struct ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpCA);
cb(ffn_inp, "ffn_inp", il);
// kq = ggml_soft_max_ext(ctx0, kq, KQ_mask_cross, 1.0f, hparams.f_max_alibi_bias);
// cb(kq, "kq_soft_max_ext", il);
// feed-forward network
{
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm, NULL,
LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
// struct ggml_tensor * v = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_reshape_2d(ctx0, Vcur, n_embd_gqa, n_outputs_enc)));
// cb(v, "v", il);
// T5 uses relu, flan-T5 uses gelu-gated
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
model.layers[il].ffn_gate, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
model.layers[il].ffn_gate_enc ? LLM_FFN_GELU : LLM_FFN_RELU,
model.layers[il].ffn_gate_enc ? LLM_FFN_PAR : LLM_FFN_SEQ,
il);
cb(cur, "ffn_out", il);
}
// struct ggml_tensor * kqv = ggml_mul_mat(ctx0, ggml_reshape_3d(ctx0, v, n_outputs_enc, n_embd_head, n_head_kv), kq);
// cb(kqv, "kqv", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cb(cur, "ffn_out", il);
// struct ggml_tensor * kqv_merged = ggml_permute(ctx0, kqv, 0, 2, 1, 3);
// cb(kqv_merged, "kqv_merged", il);
cur = lgf->build_cvec(ctx0, cur, il);
cb(cur, "l_out", il);
// cur = ggml_cont_2d(ctx0, kqv_merged, n_embd_gqa, n_tokens);
// cb(cur, "kqv_merged_cont", il);
// input for next layer
inpL = cur;
}
// ggml_build_forward_expand(gf, cur);
cur = inpL;
cb(cur, "result_embd", -1);
// cur = build_lora_mm(model.layers[il].wo_cross, cur);
// cb(cur, "kqv_out", il);
// }
cur = build_norm(cur,
model.output_norm, NULL,
LLM_NORM_RMS, -1);
// if (il == n_layer - 1) {
// // skip computing output for unused tokens
// struct ggml_tensor * inp_out_ids = build_inp_out_ids();
// cur = ggml_get_rows(ctx0, cur, inp_out_ids);
// inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
// inpCA = ggml_get_rows(ctx0, inpCA, inp_out_ids);
// }
cb(cur, "result_norm", -1);
res.t_embd = cur;
// struct ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpCA);
// cb(ffn_inp, "ffn_inp", il);
// lm_head
cur = build_lora_mm(model.output, cur);
// // feed-forward network
// {
// cur = build_norm(ffn_inp,
// model.layers[il].ffn_norm, NULL,
// LLM_NORM_RMS, il);
// cb(cur, "ffn_norm", il);
cb(cur, "result_output", -1);
res.t_logits = cur;
// // T5 uses relu, flan-T5 uses gelu-gated
// cur = build_ffn(cur,
// model.layers[il].ffn_up, NULL, NULL,
// model.layers[il].ffn_gate, NULL, NULL,
// model.layers[il].ffn_down, NULL, NULL,
// NULL,
// model.layers[il].ffn_gate_enc ? LLM_FFN_GELU : LLM_FFN_RELU,
// model.layers[il].ffn_gate_enc ? LLM_FFN_PAR : LLM_FFN_SEQ,
// il);
// cb(cur, "ffn_out", il);
// }
// cur = ggml_add(ctx0, cur, ffn_inp);
// cb(cur, "ffn_out", il);
// ggml_tensor * layer_dir = lctx.cvec.tensor_for(il);
// if (layer_dir != nullptr) {
// cur = ggml_add(ctx0, cur, layer_dir);
// }
// cb(cur, "l_out", il);
// // input for next layer
// inpL = cur;
// }
// cur = inpL;
// cb(cur, "result_embd", -1);
// cur = build_norm(cur,
// model.output_norm, NULL,
// LLM_NORM_RMS, -1);
// cb(cur, "result_norm", -1);
// res.t_embd = cur;
// // lm_head
// cur = build_lora_mm(model.output, cur);
// cb(cur, "result_output", -1);
// res.t_logits = cur;
// ggml_build_forward_expand(gf, cur);
// return gf;
//}
ggml_build_forward_expand(gf, cur);
}
void build_jais(ggml_cgraph * gf) {
const int64_t n_embd_head = hparams.n_embd_head_v;
@@ -11119,7 +11119,7 @@ llama_graph_result llama_model::build_graph(
llm.build_t5_enc(gf);
break;
case LLAMA_GRAPH_TYPE_DECODER:
//llm.build_t5_dec(gf);
llm.build_t5_dec(gf);
break;
default:
GGML_ABORT("invalid graph type");