Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3#include "llama-memory-recurrent.h"4 5#include <algorithm>6 7llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {}8 9ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,10 ggml_tensor * cur,11 const llama_model & model,12 const llama_ubatch & ubatch,13 int il) {14 const auto * mctx_cur = inp->mctx;15 16 const auto kv_head = mctx_cur->get_head();17 18 const auto & layer = model.layers[il];19 20 const int64_t d_conv = hparams.ssm_d_conv;21 const int64_t d_inner = hparams.ssm_d_inner;22 const int64_t d_state = hparams.ssm_d_state;23 const int64_t dt_rank = hparams.ssm_dt_rank;24 const int64_t n_head = d_inner;25 const int64_t head_dim = 1;26 const int64_t n_seqs = ubatch.n_seqs;27 // Some variants of Mamba arch (e.g. FalconMamba do apply layer norm on B and Dt layers)28 const bool ssm_dt_b_c_rms = hparams.ssm_dt_b_c_rms;29 30 const int64_t n_seq_tokens = ubatch.n_seq_tokens;31 32 GGML_ASSERT(n_seqs != 0);33 GGML_ASSERT(ubatch.equal_seqs());34 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);35 GGML_ASSERT(d_inner % n_head == 0);36 37 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);38 ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);39 40 ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);41 conv = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner, n_seqs);42 43 // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}44 cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);45 46 // {n_embd, 2*d_inner} @ {n_embd, n_seq_tokens, n_seqs} => {2*d_inner, n_seq_tokens, n_seqs}47 ggml_tensor * xz = build_lora_mm(layer.ssm_in, cur, layer.ssm_in_s);48 // split the above in two49 // => {d_inner, n_seq_tokens, n_seqs}50 ggml_tensor * x = ggml_view_3d(ctx0, xz, d_inner, xz->ne[1], xz->ne[2], xz->nb[1], xz->nb[2], 0);51 ggml_tensor * z =52 ggml_view_3d(ctx0, xz, d_inner, xz->ne[1], xz->ne[2], xz->nb[1], xz->nb[2], d_inner * ggml_element_size(xz));53 54 // conv55 {56 // => {d_conv - 1 + n_seq_tokens, d_inner, n_seqs}57 ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, x), 0);58 59 // copy last (d_conv - 1) columns back into the state cache60 ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, conv_x->nb[1], conv_x->nb[2],61 n_seq_tokens * (conv_x->nb[0]));62 63 ggml_build_forward_expand(64 gf, ggml_cpy(ctx0, last_conv,65 ggml_view_1d(ctx0, conv_states_all, (d_conv - 1) * (d_inner) * (n_seqs),66 kv_head * (d_conv - 1) * (d_inner) *ggml_element_size(conv_states_all))));67 68 // 1D convolution69 // The equivalent is to make a self-overlapping view of conv_x70 // over d_conv columns at each stride in the 3rd dimension,71 // then element-wise multiply that with the conv1d weight,72 // then sum the elements of each row,73 // (the last two steps are a dot product over rows (also doable with mul_mat))74 // then permute away the ne[0] dimension,75 // and then you're left with the resulting x tensor.76 // For simultaneous sequences, all sequences need to have the same length.77 x = ggml_ssm_conv(ctx0, conv_x, layer.ssm_conv1d);78 79 // bias80 x = ggml_add(ctx0, x, layer.ssm_conv1d_b);81 82 x = ggml_silu(ctx0, x);83 }84 85 // ssm86 {87 // {d_inner, dt_rank + 2*d_state} @ {d_inner, n_seq_tokens, n_seqs} => {dt_rank + 2*d_state, n_seq_tokens, n_seqs}88 ggml_tensor * x_db = build_lora_mm(layer.ssm_x, x);89 // split90 ggml_tensor * dt = ggml_view_3d(ctx0, x_db, dt_rank, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], 0);91 ggml_tensor * B =92 ggml_view_4d(ctx0, x_db, d_state, /* n_group */ 1, n_seq_tokens, n_seqs, d_state * x_db->nb[0], x_db->nb[1],93 x_db->nb[2], ggml_element_size(x_db) * dt_rank);94 ggml_tensor * C =95 ggml_view_4d(ctx0, x_db, d_state, /* n_group */ 1, n_seq_tokens, n_seqs, d_state * x_db->nb[0], x_db->nb[1],96 x_db->nb[2], ggml_element_size(x_db) * (dt_rank + d_state));97 98 // Some Mamba variants (e.g. FalconMamba, Jamba) apply RMS norm in B, C & Dt layers99 if (ssm_dt_b_c_rms || (layer.ssm_dt_norm && layer.ssm_b_norm && layer.ssm_c_norm)) {100 dt = build_norm(dt, layer.ssm_dt_norm, NULL, LLM_NORM_RMS, il);101 B = build_norm(B, layer.ssm_b_norm, NULL, LLM_NORM_RMS, il);102 C = build_norm(C, layer.ssm_c_norm, NULL, LLM_NORM_RMS, il);103 }104 105 // {dt_rank, d_inner} @ {dt_rank, n_seq_tokens, n_seqs} => {d_inner, n_seq_tokens, n_seqs}106 dt = build_lora_mm(layer.ssm_dt, dt);107 dt = ggml_add(ctx0, dt, layer.ssm_dt_b);108 109 cur = x;110 x = ggml_reshape_4d(ctx0, x, head_dim, n_head, n_seq_tokens, n_seqs);111 112 ggml_tensor * A = layer.ssm_a;113 114 // use the states and the indices provided by build_recurrent_state115 // (this is necessary in order to properly use the states before they are overwritten,116 // while avoiding to make unnecessary copies of the states)117 auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {118 ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size());119 120 // Custom operator to optimize the parallel associative scan121 // as described in the Annex D of the Mamba paper.122 // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}123 return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);124 };125 126 ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);127 128 // store last states129 ggml_build_forward_expand(130 gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, x->nb[3] * x->ne[3]),131 ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs,132 kv_head * d_state * d_inner * ggml_element_size(ssm_states_all))));133 134 ggml_tensor * y = ggml_view_3d(ctx0, y_ssm, d_inner, n_seq_tokens, n_seqs, x->nb[2], x->nb[3], 0);135 136 // TODO: skip computing output earlier for unused tokens137 138 y = ggml_add(ctx0, y, ggml_mul(ctx0, cur, layer.ssm_d));139 y = ggml_swiglu_split(ctx0, ggml_cont(ctx0, z), y);140 141 // {d_inner, n_embd} @ {d_inner, n_seq_tokens, n_seqs} => {n_embd, n_seq_tokens, n_seqs}142 cur = build_lora_mm(layer.ssm_out, y, layer.ssm_out_s);143 }144 145 // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}146 cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens * n_seqs);147 148 return cur;149}150 151ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,152 ggml_tensor * cur,153 const llama_model & model,154 const llama_ubatch & ubatch,155 int il) const {156 const auto * mctx_cur = inp->mctx;157 158 const auto kv_head = mctx_cur->get_head();159 const auto mem_size = mctx_cur->get_size();160 161 const int64_t d_conv = hparams.ssm_d_conv;162 const int64_t d_inner = hparams.ssm_d_inner;163 const int64_t d_state = hparams.ssm_d_state;164 const int64_t n_head = hparams.ssm_dt_rank;165 const int64_t head_dim = d_inner / n_head;166 const int64_t n_group = hparams.ssm_n_group;167 const int64_t n_seqs = ubatch.n_seqs;168 169 const int64_t n_seq_tokens = ubatch.n_seq_tokens;170 const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;171 172 GGML_ASSERT(n_seqs != 0);173 GGML_ASSERT(ubatch.equal_seqs());174 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);175 GGML_ASSERT(d_inner % n_head == 0);176 GGML_ASSERT(d_inner % n_group == 0);177 178 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);179 ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);180 const int64_t state_slots = ssm_states_all->ne[1];181 182 ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);183 conv = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);184 185 // d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads186 187 // Keep the projection 2D: with a {n_embd, 1, n_seqs} batch the CUDA backend188 // dispatches a column-batched GEMV for what is a large dense GEMM.189 // {n_embd, d_in_proj} @ {n_embd, n_tokens} => {d_in_proj, n_tokens}190 ggml_tensor * zxBCdt = build_lora_mm(model.layers[il].ssm_in, cur, model.layers[il].ssm_in_s);191 // {d_in_proj, n_tokens} => {d_in_proj, n_seq_tokens, n_seqs}192 zxBCdt = ggml_reshape_3d(ctx0, zxBCdt, zxBCdt->ne[0], n_seq_tokens, n_seqs);193 194 // split the above in three195 ggml_tensor * z = ggml_view_4d(ctx0, zxBCdt, head_dim, n_head, n_seq_tokens, n_seqs, head_dim * zxBCdt->nb[0],196 zxBCdt->nb[1], zxBCdt->nb[2], 0);197 ggml_tensor * xBC = ggml_view_3d(ctx0, zxBCdt, d_inner + 2 * n_group * d_state, n_seq_tokens, n_seqs, zxBCdt->nb[1],198 zxBCdt->nb[2], d_inner * ggml_element_size(zxBCdt));199 ggml_tensor * dt = ggml_view_3d(ctx0, zxBCdt, n_head, n_seq_tokens, n_seqs, zxBCdt->nb[1], zxBCdt->nb[2],200 (2 * d_inner + 2 * n_group * d_state) * ggml_element_size(zxBCdt));201 202 // conv203 {204 // => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}205 ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);206 207 const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);208 const size_t row_size = ggml_row_size(conv_states_all->type, row_count);209 const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);210 211 for (int64_t slot = 0; slot < n_written; ++slot) {212 ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,213 conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);214 215 ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,216 ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,217 conv_states_all->nb[1],218 ((size_t) slot * mem_size + kv_head) * row_size)));219 }220 221 // 1D convolution222 // The equivalent is to make a self-overlapping view of conv_x223 // over d_conv columns at each stride in the 3rd dimension,224 // then element-wise multiply that with the conv1d weight,225 // then sum the elements of each row,226 // (the last two steps are a dot product over rows (also doable with mul_mat))227 // then permute away the ne[0] dimension,228 // and then you're left with the resulting x tensor.229 // For simultaneous sequences, all sequences need to have the same length.230 xBC = ggml_ssm_conv(ctx0, conv_x, model.layers[il].ssm_conv1d);231 232 // bias233 xBC = ggml_add(ctx0, xBC, model.layers[il].ssm_conv1d_b);234 235 xBC = ggml_silu(ctx0, xBC);236 }237 238 // ssm239 {240 // These correspond to V K Q in SSM/attention duality241 ggml_tensor * x = ggml_view_4d(ctx0, xBC, head_dim, n_head, n_seq_tokens, n_seqs, head_dim * xBC->nb[0],242 xBC->nb[1], xBC->nb[2], 0);243 ggml_tensor * B = ggml_view_4d(ctx0, xBC, d_state, n_group, n_seq_tokens, n_seqs, d_state * xBC->nb[0],244 xBC->nb[1], xBC->nb[2], d_inner * ggml_element_size(xBC));245 ggml_tensor * C = ggml_view_4d(ctx0, xBC, d_state, n_group, n_seq_tokens, n_seqs, d_state * xBC->nb[0],246 xBC->nb[1], xBC->nb[2], (d_inner + n_group * d_state) * ggml_element_size(xBC));247 248 // {n_head, n_seq_tokens, n_seqs}249 dt = ggml_add(ctx0, ggml_cont(ctx0, dt), model.layers[il].ssm_dt_b);250 251 ggml_tensor * A = model.layers[il].ssm_a;252 253 // use the states and the indices provided by build_recurrent_state254 // (this is necessary in order to properly use the states before they are overwritten,255 // while avoiding to make unnecessary copies of the states)256 auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {257 ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots);258 259 // TODO: use semistructured matrices to implement state-space duality260 // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}261 // K > 1 asks the backend to return rollback snapshots in addition to the final state.262 return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K);263 };264 265 ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);266 const int64_t D = d_state * d_inner;267 const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);268 const size_t row_size = ggml_row_size(ssm_states_all->type, D);269 const size_t y_row_size = ggml_row_size(y_ssm->type, D);270 const size_t state_offset = ggml_nelements(x) * ggml_element_size(x);271 272 ggml_build_forward_expand(273 gf, ggml_cpy(ctx0,274 ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written,275 y_row_size, y_row_size * n_seqs, state_offset),276 ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written,277 ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size)));278 279 ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1],280 n_seq_tokens * n_head * x->nb[1], 0);281 282 // TODO: skip computing output earlier for unused tokens283 284 y = ggml_add(ctx0, y, ggml_mul(ctx0, x, model.layers[il].ssm_d));285 cb(y, "mamba2_y_add_d", il);286 y = ggml_swiglu_split(ctx0, ggml_cont(ctx0, z), y);287 288 // grouped RMS norm289 if (model.layers[il].ssm_norm) {290 y = ggml_reshape_4d(ctx0, y, d_inner / n_group, n_group, n_seq_tokens, n_seqs);291 y = build_norm(y, model.layers[il].ssm_norm, NULL, LLM_NORM_RMS, il);292 }293 294 y = ggml_reshape_2d(ctx0, y, d_inner, n_seq_tokens * n_seqs);295 296 // {d_inner, n_embd} @ {d_inner, n_tokens} => {n_embd, n_tokens}297 cur = build_lora_mm(model.layers[il].ssm_out, y, model.layers[il].ssm_out_s);298 }299 300 cb(cur, "mamba_out", il);301 return cur;302}303